grafeo_engine/database/

persistence.rs

//! Persistence, snapshots, and data export for GrafeoDB.

#[cfg(feature = "wal")]
use std::path::Path;

use grafeo_common::types::{EdgeId, NodeId, Value};
use grafeo_common::utils::error::{Error, Result};
use hashbrown::HashSet;

use crate::config::Config;

#[cfg(feature = "wal")]
use grafeo_adapters::storage::wal::WalRecord;

use crate::catalog::{
    EdgeTypeDefinition, GraphTypeDefinition, NodeTypeDefinition, ProcedureDefinition,
};

/// Current snapshot version.
const SNAPSHOT_VERSION: u8 = 3;

/// Binary snapshot format (v3: graph data, named graphs, RDF, and schema).
#[derive(serde::Serialize, serde::Deserialize)]
struct Snapshot {
    version: u8,
    nodes: Vec<SnapshotNode>,
    edges: Vec<SnapshotEdge>,
    named_graphs: Vec<NamedGraphSnapshot>,
    rdf_triples: Vec<SnapshotTriple>,
    rdf_named_graphs: Vec<RdfNamedGraphSnapshot>,
    schema: SnapshotSchema,
}

/// Schema metadata within a snapshot.
#[derive(serde::Serialize, serde::Deserialize, Default)]
struct SnapshotSchema {
    node_types: Vec<NodeTypeDefinition>,
    edge_types: Vec<EdgeTypeDefinition>,
    graph_types: Vec<GraphTypeDefinition>,
    procedures: Vec<ProcedureDefinition>,
    schemas: Vec<String>,
    graph_type_bindings: Vec<(String, String)>,
}

/// A named graph partition within a v2 snapshot.
#[derive(serde::Serialize, serde::Deserialize)]
struct NamedGraphSnapshot {
    name: String,
    nodes: Vec<SnapshotNode>,
    edges: Vec<SnapshotEdge>,
}

/// An RDF triple in snapshot format (N-Triples encoded terms).
#[derive(serde::Serialize, serde::Deserialize)]
struct SnapshotTriple {
    subject: String,
    predicate: String,
    object: String,
}

/// An RDF named graph in snapshot format.
#[derive(serde::Serialize, serde::Deserialize)]
struct RdfNamedGraphSnapshot {
    name: String,
    triples: Vec<SnapshotTriple>,
}

#[derive(serde::Serialize, serde::Deserialize)]
struct SnapshotNode {
    id: NodeId,
    labels: Vec<String>,
    properties: Vec<(String, Value)>,
}

#[derive(serde::Serialize, serde::Deserialize)]
struct SnapshotEdge {
    id: EdgeId,
    src: NodeId,
    dst: NodeId,
    edge_type: String,
    properties: Vec<(String, Value)>,
}

/// Collects all nodes from a store into snapshot format.
fn collect_snapshot_nodes(store: &grafeo_core::graph::lpg::LpgStore) -> Vec<SnapshotNode> {
    store
        .all_nodes()
        .map(|n| SnapshotNode {
            id: n.id,
            labels: n.labels.iter().map(|l| l.to_string()).collect(),
            properties: n
                .properties
                .into_iter()
                .map(|(k, v)| (k.to_string(), v))
                .collect(),
        })
        .collect()
}

/// Collects all edges from a store into snapshot format.
fn collect_snapshot_edges(store: &grafeo_core::graph::lpg::LpgStore) -> Vec<SnapshotEdge> {
    store
        .all_edges()
        .map(|e| SnapshotEdge {
            id: e.id,
            src: e.src,
            dst: e.dst,
            edge_type: e.edge_type.to_string(),
            properties: e
                .properties
                .into_iter()
                .map(|(k, v)| (k.to_string(), v))
                .collect(),
        })
        .collect()
}

/// Populates a store from snapshot node/edge data.
fn populate_store_from_snapshot(
    store: &grafeo_core::graph::lpg::LpgStore,
    nodes: Vec<SnapshotNode>,
    edges: Vec<SnapshotEdge>,
) -> Result<()> {
    for node in nodes {
        let label_refs: Vec<&str> = node.labels.iter().map(|s| s.as_str()).collect();
        store.create_node_with_id(node.id, &label_refs)?;
        for (key, value) in node.properties {
            store.set_node_property(node.id, &key, value);
        }
    }
    for edge in edges {
        store.create_edge_with_id(edge.id, edge.src, edge.dst, &edge.edge_type)?;
        for (key, value) in edge.properties {
            store.set_edge_property(edge.id, &key, value);
        }
    }
    Ok(())
}

/// Validates snapshot nodes/edges for duplicates and dangling references.
fn validate_snapshot_data(nodes: &[SnapshotNode], edges: &[SnapshotEdge]) -> Result<()> {
    let mut node_ids = HashSet::with_capacity(nodes.len());
    for node in nodes {
        if !node_ids.insert(node.id) {
            return Err(Error::Internal(format!(
                "snapshot contains duplicate node ID {}",
                node.id
            )));
        }
    }
    let mut edge_ids = HashSet::with_capacity(edges.len());
    for edge in edges {
        if !edge_ids.insert(edge.id) {
            return Err(Error::Internal(format!(
                "snapshot contains duplicate edge ID {}",
                edge.id
            )));
        }
        if !node_ids.contains(&edge.src) {
            return Err(Error::Internal(format!(
                "snapshot edge {} references non-existent source node {}",
                edge.id, edge.src
            )));
        }
        if !node_ids.contains(&edge.dst) {
            return Err(Error::Internal(format!(
                "snapshot edge {} references non-existent destination node {}",
                edge.id, edge.dst
            )));
        }
    }
    Ok(())
}

/// Collects all triples from an RDF store into snapshot format.
#[cfg(feature = "rdf")]
fn collect_rdf_triples(store: &grafeo_core::graph::rdf::RdfStore) -> Vec<SnapshotTriple> {
    store
        .triples()
        .into_iter()
        .map(|t| SnapshotTriple {
            subject: t.subject().to_string(),
            predicate: t.predicate().to_string(),
            object: t.object().to_string(),
        })
        .collect()
}

/// Populates an RDF store from snapshot triples.
#[cfg(feature = "rdf")]
fn populate_rdf_store(store: &grafeo_core::graph::rdf::RdfStore, triples: &[SnapshotTriple]) {
    use grafeo_core::graph::rdf::{Term, Triple};
    for triple in triples {
        if let (Some(s), Some(p), Some(o)) = (
            Term::from_ntriples(&triple.subject),
            Term::from_ntriples(&triple.predicate),
            Term::from_ntriples(&triple.object),
        ) {
            store.insert(Triple::new(s, p, o));
        }
    }
}

// =========================================================================
// Snapshot deserialization helpers (used by single-file format)
// =========================================================================

/// Decodes snapshot bytes and populates a store and catalog.
#[cfg(feature = "grafeo-file")]
pub(super) fn load_snapshot_into_store(
    store: &std::sync::Arc<grafeo_core::graph::lpg::LpgStore>,
    catalog: &std::sync::Arc<crate::catalog::Catalog>,
    #[cfg(feature = "rdf")] rdf_store: &std::sync::Arc<grafeo_core::graph::rdf::RdfStore>,
    data: &[u8],
) -> grafeo_common::utils::error::Result<()> {
    use grafeo_common::utils::error::Error;

    let config = bincode::config::standard();
    let (snapshot, _) =
        bincode::serde::decode_from_slice::<Snapshot, _>(data, config).map_err(|e| {
            Error::Serialization(format!("failed to decode snapshot from .grafeo file: {e}"))
        })?;

    populate_store_from_snapshot_ref(store, &snapshot.nodes, &snapshot.edges)?;
    for graph in &snapshot.named_graphs {
        store
            .create_graph(&graph.name)
            .map_err(|e| Error::Internal(e.to_string()))?;
        if let Some(graph_store) = store.graph(&graph.name) {
            populate_store_from_snapshot_ref(&graph_store, &graph.nodes, &graph.edges)?;
        }
    }
    restore_schema_from_snapshot(catalog, &snapshot.schema);

    // Restore RDF triples
    #[cfg(feature = "rdf")]
    {
        populate_rdf_store(rdf_store, &snapshot.rdf_triples);
        for rdf_graph in &snapshot.rdf_named_graphs {
            rdf_store.create_graph(&rdf_graph.name);
            if let Some(graph_store) = rdf_store.graph(&rdf_graph.name) {
                populate_rdf_store(&graph_store, &rdf_graph.triples);
            }
        }
    }

    Ok(())
}

/// Populates a store from snapshot refs (borrowed, for single-file loading).
#[cfg(feature = "grafeo-file")]
fn populate_store_from_snapshot_ref(
    store: &grafeo_core::graph::lpg::LpgStore,
    nodes: &[SnapshotNode],
    edges: &[SnapshotEdge],
) -> grafeo_common::utils::error::Result<()> {
    for node in nodes {
        let label_refs: Vec<&str> = node.labels.iter().map(|s| s.as_str()).collect();
        store.create_node_with_id(node.id, &label_refs)?;
        for (key, value) in &node.properties {
            store.set_node_property(node.id, key, value.clone());
        }
    }
    for edge in edges {
        store.create_edge_with_id(edge.id, edge.src, edge.dst, &edge.edge_type)?;
        for (key, value) in &edge.properties {
            store.set_edge_property(edge.id, key, value.clone());
        }
    }
    Ok(())
}

/// Restores schema definitions from a snapshot into the catalog.
fn restore_schema_from_snapshot(
    catalog: &std::sync::Arc<crate::catalog::Catalog>,
    schema: &SnapshotSchema,
) {
    for def in &schema.node_types {
        catalog.register_or_replace_node_type(def.clone());
    }
    for def in &schema.edge_types {
        catalog.register_or_replace_edge_type_def(def.clone());
    }
    for def in &schema.graph_types {
        let _ = catalog.register_graph_type(def.clone());
    }
    for def in &schema.procedures {
        catalog.replace_procedure(def.clone()).ok();
    }
    for name in &schema.schemas {
        let _ = catalog.register_schema_namespace(name.clone());
    }
    for (graph_name, type_name) in &schema.graph_type_bindings {
        let _ = catalog.bind_graph_type(graph_name, type_name.clone());
    }
}

/// Collects schema definitions from the catalog into snapshot format.
fn collect_schema(catalog: &std::sync::Arc<crate::catalog::Catalog>) -> SnapshotSchema {
    SnapshotSchema {
        node_types: catalog.all_node_type_defs(),
        edge_types: catalog.all_edge_type_defs(),
        graph_types: catalog.all_graph_type_defs(),
        procedures: catalog.all_procedure_defs(),
        schemas: catalog.schema_names(),
        graph_type_bindings: catalog.all_graph_type_bindings(),
    }
}

impl super::GrafeoDB {
    // =========================================================================
    // ADMIN API: Persistence Control
    // =========================================================================

    /// Saves the database to a file path.
    ///
    /// - If the path ends in `.grafeo`: creates a single-file database
    /// - Otherwise: creates a WAL directory-backed database at the path
    /// - If in-memory: creates a new persistent database at path
    /// - If file-backed: creates a copy at the new path
    ///
    /// The original database remains unchanged.
    ///
    /// # Errors
    ///
    /// Returns an error if the save operation fails.
    ///
    /// Requires the `wal` feature for persistence support.
    #[cfg(feature = "wal")]
    pub fn save(&self, path: impl AsRef<Path>) -> Result<()> {
        let path = path.as_ref();

        // Single-file format: export snapshot directly to a .grafeo file
        #[cfg(feature = "grafeo-file")]
        if path.extension().is_some_and(|ext| ext == "grafeo") {
            return self.save_as_grafeo_file(path);
        }

        // Create target database with WAL enabled
        let target_config = Config::persistent(path);
        let target = Self::with_config(target_config)?;

        // Copy all nodes using WAL-enabled methods
        for node in self.store.all_nodes() {
            let label_refs: Vec<&str> = node.labels.iter().map(|s| &**s).collect();
            target.store.create_node_with_id(node.id, &label_refs)?;

            // Log to WAL
            target.log_wal(&WalRecord::CreateNode {
                id: node.id,
                labels: node.labels.iter().map(|s| s.to_string()).collect(),
            })?;

            // Copy properties
            for (key, value) in node.properties {
                target
                    .store
                    .set_node_property(node.id, key.as_str(), value.clone());
                target.log_wal(&WalRecord::SetNodeProperty {
                    id: node.id,
                    key: key.to_string(),
                    value,
                })?;
            }
        }

        // Copy all edges using WAL-enabled methods
        for edge in self.store.all_edges() {
            target
                .store
                .create_edge_with_id(edge.id, edge.src, edge.dst, &edge.edge_type)?;

            // Log to WAL
            target.log_wal(&WalRecord::CreateEdge {
                id: edge.id,
                src: edge.src,
                dst: edge.dst,
                edge_type: edge.edge_type.to_string(),
            })?;

            // Copy properties
            for (key, value) in edge.properties {
                target
                    .store
                    .set_edge_property(edge.id, key.as_str(), value.clone());
                target.log_wal(&WalRecord::SetEdgeProperty {
                    id: edge.id,
                    key: key.to_string(),
                    value,
                })?;
            }
        }

        // Copy named graphs
        for graph_name in self.store.graph_names() {
            if let Some(src_graph) = self.store.graph(&graph_name) {
                target.log_wal(&WalRecord::CreateNamedGraph {
                    name: graph_name.clone(),
                })?;
                target
                    .store
                    .create_graph(&graph_name)
                    .map_err(|e| Error::Internal(e.to_string()))?;

                if let Some(dst_graph) = target.store.graph(&graph_name) {
                    // Switch WAL context to this named graph
                    target.log_wal(&WalRecord::SwitchGraph {
                        name: Some(graph_name.clone()),
                    })?;

                    for node in src_graph.all_nodes() {
                        let label_refs: Vec<&str> = node.labels.iter().map(|s| &**s).collect();
                        dst_graph.create_node_with_id(node.id, &label_refs)?;
                        target.log_wal(&WalRecord::CreateNode {
                            id: node.id,
                            labels: node.labels.iter().map(|s| s.to_string()).collect(),
                        })?;
                        for (key, value) in node.properties {
                            dst_graph.set_node_property(node.id, key.as_str(), value.clone());
                            target.log_wal(&WalRecord::SetNodeProperty {
                                id: node.id,
                                key: key.to_string(),
                                value,
                            })?;
                        }
                    }
                    for edge in src_graph.all_edges() {
                        dst_graph.create_edge_with_id(
                            edge.id,
                            edge.src,
                            edge.dst,
                            &edge.edge_type,
                        )?;
                        target.log_wal(&WalRecord::CreateEdge {
                            id: edge.id,
                            src: edge.src,
                            dst: edge.dst,
                            edge_type: edge.edge_type.to_string(),
                        })?;
                        for (key, value) in edge.properties {
                            dst_graph.set_edge_property(edge.id, key.as_str(), value.clone());
                            target.log_wal(&WalRecord::SetEdgeProperty {
                                id: edge.id,
                                key: key.to_string(),
                                value,
                            })?;
                        }
                    }
                }
            }
        }

        // Switch WAL context back to default graph
        if !self.store.graph_names().is_empty() {
            target.log_wal(&WalRecord::SwitchGraph { name: None })?;
        }

        // Copy RDF data with WAL logging
        #[cfg(feature = "rdf")]
        {
            for triple in self.rdf_store.triples() {
                let record = WalRecord::InsertRdfTriple {
                    subject: triple.subject().to_string(),
                    predicate: triple.predicate().to_string(),
                    object: triple.object().to_string(),
                    graph: None,
                };
                target.rdf_store.insert((*triple).clone());
                target.log_wal(&record)?;
            }
            for name in self.rdf_store.graph_names() {
                target.log_wal(&WalRecord::CreateRdfGraph { name: name.clone() })?;
                if let Some(src_graph) = self.rdf_store.graph(&name) {
                    let dst_graph = target.rdf_store.graph_or_create(&name);
                    for triple in src_graph.triples() {
                        let record = WalRecord::InsertRdfTriple {
                            subject: triple.subject().to_string(),
                            predicate: triple.predicate().to_string(),
                            object: triple.object().to_string(),
                            graph: Some(name.clone()),
                        };
                        dst_graph.insert((*triple).clone());
                        target.log_wal(&record)?;
                    }
                }
            }
        }

        // Checkpoint and close the target database
        target.close()?;

        Ok(())
    }

    /// Creates an in-memory copy of this database.
    ///
    /// Returns a new database that is completely independent, including
    /// all named graph data.
    /// Useful for:
    /// Saves the database to a single `.grafeo` file.
    #[cfg(feature = "grafeo-file")]
    fn save_as_grafeo_file(&self, path: &Path) -> Result<()> {
        use grafeo_adapters::storage::file::GrafeoFileManager;

        let snapshot_data = self.export_snapshot()?;
        let epoch = self.store.current_epoch();
        let transaction_id = self
            .transaction_manager
            .last_assigned_transaction_id()
            .map_or(0, |t| t.0);
        let node_count = self.store.node_count() as u64;
        let edge_count = self.store.edge_count() as u64;

        let fm = GrafeoFileManager::create(path)?;
        fm.write_snapshot(
            &snapshot_data,
            epoch.0,
            transaction_id,
            node_count,
            edge_count,
        )?;
        Ok(())
    }

    /// - Testing modifications without affecting the original
    /// - Faster operations when persistence isn't needed
    ///
    /// # Errors
    ///
    /// Returns an error if the copy operation fails.
    pub fn to_memory(&self) -> Result<Self> {
        let config = Config::in_memory();
        let target = Self::with_config(config)?;

        // Copy default graph nodes
        for node in self.store.all_nodes() {
            let label_refs: Vec<&str> = node.labels.iter().map(|s| &**s).collect();
            target.store.create_node_with_id(node.id, &label_refs)?;
            for (key, value) in node.properties {
                target.store.set_node_property(node.id, key.as_str(), value);
            }
        }

        // Copy default graph edges
        for edge in self.store.all_edges() {
            target
                .store
                .create_edge_with_id(edge.id, edge.src, edge.dst, &edge.edge_type)?;
            for (key, value) in edge.properties {
                target.store.set_edge_property(edge.id, key.as_str(), value);
            }
        }

        // Copy named graphs
        for graph_name in self.store.graph_names() {
            if let Some(src_graph) = self.store.graph(&graph_name) {
                target
                    .store
                    .create_graph(&graph_name)
                    .map_err(|e| Error::Internal(e.to_string()))?;
                if let Some(dst_graph) = target.store.graph(&graph_name) {
                    for node in src_graph.all_nodes() {
                        let label_refs: Vec<&str> = node.labels.iter().map(|s| &**s).collect();
                        dst_graph.create_node_with_id(node.id, &label_refs)?;
                        for (key, value) in node.properties {
                            dst_graph.set_node_property(node.id, key.as_str(), value);
                        }
                    }
                    for edge in src_graph.all_edges() {
                        dst_graph.create_edge_with_id(
                            edge.id,
                            edge.src,
                            edge.dst,
                            &edge.edge_type,
                        )?;
                        for (key, value) in edge.properties {
                            dst_graph.set_edge_property(edge.id, key.as_str(), value);
                        }
                    }
                }
            }
        }

        // Copy RDF data
        #[cfg(feature = "rdf")]
        {
            for triple in self.rdf_store.triples() {
                target.rdf_store.insert((*triple).clone());
            }
            for name in self.rdf_store.graph_names() {
                if let Some(src_graph) = self.rdf_store.graph(&name) {
                    let dst_graph = target.rdf_store.graph_or_create(&name);
                    for triple in src_graph.triples() {
                        dst_graph.insert((*triple).clone());
                    }
                }
            }
        }

        Ok(target)
    }

    /// Opens a database file and loads it entirely into memory.
    ///
    /// The returned database has no connection to the original file.
    /// Changes will NOT be written back to the file.
    ///
    /// # Errors
    ///
    /// Returns an error if the file can't be opened or loaded.
    #[cfg(feature = "wal")]
    pub fn open_in_memory(path: impl AsRef<Path>) -> Result<Self> {
        // Open the source database (triggers WAL recovery)
        let source = Self::open(path)?;

        // Create in-memory copy
        let target = source.to_memory()?;

        // Close the source (releases file handles)
        source.close()?;

        Ok(target)
    }

    // =========================================================================
    // ADMIN API: Snapshot Export/Import
    // =========================================================================

    /// Exports the entire database to a binary snapshot (v2 format).
    ///
    /// The returned bytes can be stored (e.g. in IndexedDB) and later
    /// restored with [`import_snapshot()`](Self::import_snapshot).
    /// Includes all named graph data.
    ///
    /// # Errors
    ///
    /// Returns an error if serialization fails.
    pub fn export_snapshot(&self) -> Result<Vec<u8>> {
        let nodes = collect_snapshot_nodes(&self.store);
        let edges = collect_snapshot_edges(&self.store);

        // Collect named graphs
        let named_graphs: Vec<NamedGraphSnapshot> = self
            .store
            .graph_names()
            .into_iter()
            .filter_map(|name| {
                self.store
                    .graph(&name)
                    .map(|graph_store| NamedGraphSnapshot {
                        name,
                        nodes: collect_snapshot_nodes(&graph_store),
                        edges: collect_snapshot_edges(&graph_store),
                    })
            })
            .collect();

        // Collect RDF triples
        #[cfg(feature = "rdf")]
        let rdf_triples = collect_rdf_triples(&self.rdf_store);
        #[cfg(not(feature = "rdf"))]
        let rdf_triples = Vec::new();

        #[cfg(feature = "rdf")]
        let rdf_named_graphs: Vec<RdfNamedGraphSnapshot> = self
            .rdf_store
            .graph_names()
            .into_iter()
            .filter_map(|name| {
                self.rdf_store
                    .graph(&name)
                    .map(|graph| RdfNamedGraphSnapshot {
                        name,
                        triples: collect_rdf_triples(&graph),
                    })
            })
            .collect();
        #[cfg(not(feature = "rdf"))]
        let rdf_named_graphs = Vec::new();

        let schema = collect_schema(&self.catalog);

        let snapshot = Snapshot {
            version: SNAPSHOT_VERSION,
            nodes,
            edges,
            named_graphs,
            rdf_triples,
            rdf_named_graphs,
            schema,
        };

        let config = bincode::config::standard();
        bincode::serde::encode_to_vec(&snapshot, config)
            .map_err(|e| Error::Internal(format!("snapshot export failed: {e}")))
    }

    /// Creates a new in-memory database from a binary snapshot.
    ///
    /// Accepts both v1 (no named graphs) and v2 (with named graphs) formats.
    /// The `data` must have been produced by [`export_snapshot()`](Self::export_snapshot).
    ///
    /// All edge references are validated before any data is inserted: every
    /// edge's source and destination must reference a node present in the
    /// snapshot, and duplicate node/edge IDs are rejected. If validation
    /// fails, no database is created.
    ///
    /// # Errors
    ///
    /// Returns an error if the snapshot is invalid, contains dangling edge
    /// references, has duplicate IDs, or deserialization fails.
    pub fn import_snapshot(data: &[u8]) -> Result<Self> {
        if data.is_empty() {
            return Err(Error::Internal("empty snapshot data".to_string()));
        }

        // Peek at version byte (bincode standard encodes u8 as raw byte)
        if data[0] != SNAPSHOT_VERSION {
            return Err(Error::Internal(format!(
                "unsupported snapshot version: {} (expected {SNAPSHOT_VERSION})",
                data[0]
            )));
        }

        let config = bincode::config::standard();
        let (snapshot, _): (Snapshot, _) = bincode::serde::decode_from_slice(data, config)
            .map_err(|e| Error::Internal(format!("snapshot import failed: {e}")))?;

        // Validate default graph data
        validate_snapshot_data(&snapshot.nodes, &snapshot.edges)?;

        // Validate each named graph
        for ng in &snapshot.named_graphs {
            validate_snapshot_data(&ng.nodes, &ng.edges)?;
        }

        let db = Self::new_in_memory();
        populate_store_from_snapshot(&db.store, snapshot.nodes, snapshot.edges)?;

        // Restore named graphs
        for ng in snapshot.named_graphs {
            db.store
                .create_graph(&ng.name)
                .map_err(|e| Error::Internal(e.to_string()))?;
            if let Some(graph_store) = db.store.graph(&ng.name) {
                populate_store_from_snapshot(&graph_store, ng.nodes, ng.edges)?;
            }
        }

        // Restore RDF triples
        #[cfg(feature = "rdf")]
        {
            populate_rdf_store(&db.rdf_store, &snapshot.rdf_triples);
            for rng in &snapshot.rdf_named_graphs {
                let graph = db.rdf_store.graph_or_create(&rng.name);
                populate_rdf_store(&graph, &rng.triples);
            }
        }

        // Restore schema
        restore_schema_from_snapshot(&db.catalog, &snapshot.schema);

        Ok(db)
    }

    /// Replaces the current database contents with data from a binary snapshot.
    ///
    /// Accepts both v1 and v2 snapshot formats. The `data` must have been
    /// produced by [`export_snapshot()`](Self::export_snapshot).
    ///
    /// All validation (duplicate IDs, dangling edge references) is performed
    /// before any data is modified. If validation fails, the current database
    /// is left unchanged. If validation passes, the store is cleared and
    /// rebuilt from the snapshot atomically (from the perspective of
    /// subsequent queries).
    ///
    /// # Errors
    ///
    /// Returns an error if the snapshot is invalid, contains dangling edge
    /// references, has duplicate IDs, or deserialization fails.
    pub fn restore_snapshot(&self, data: &[u8]) -> Result<()> {
        if data.is_empty() {
            return Err(Error::Internal("empty snapshot data".to_string()));
        }

        if data[0] != SNAPSHOT_VERSION {
            return Err(Error::Internal(format!(
                "unsupported snapshot version: {} (expected {SNAPSHOT_VERSION})",
                data[0]
            )));
        }

        let config = bincode::config::standard();
        let (snapshot, _): (Snapshot, _) = bincode::serde::decode_from_slice(data, config)
            .map_err(|e| Error::Internal(format!("snapshot restore failed: {e}")))?;

        // Validate all data before making any changes
        validate_snapshot_data(&snapshot.nodes, &snapshot.edges)?;
        for ng in &snapshot.named_graphs {
            validate_snapshot_data(&ng.nodes, &ng.edges)?;
        }

        // Drop all existing named graphs, then clear default store
        for name in self.store.graph_names() {
            self.store.drop_graph(&name);
        }
        self.store.clear();

        populate_store_from_snapshot(&self.store, snapshot.nodes, snapshot.edges)?;

        // Restore named graphs
        for ng in snapshot.named_graphs {
            self.store
                .create_graph(&ng.name)
                .map_err(|e| Error::Internal(e.to_string()))?;
            if let Some(graph_store) = self.store.graph(&ng.name) {
                populate_store_from_snapshot(&graph_store, ng.nodes, ng.edges)?;
            }
        }

        // Restore RDF data
        #[cfg(feature = "rdf")]
        {
            // Clear existing RDF data
            self.rdf_store.clear();
            for name in self.rdf_store.graph_names() {
                self.rdf_store.drop_graph(&name);
            }
            populate_rdf_store(&self.rdf_store, &snapshot.rdf_triples);
            for rng in &snapshot.rdf_named_graphs {
                let graph = self.rdf_store.graph_or_create(&rng.name);
                populate_rdf_store(&graph, &rng.triples);
            }
        }

        // Restore schema
        restore_schema_from_snapshot(&self.catalog, &snapshot.schema);

        Ok(())
    }

    // =========================================================================
    // ADMIN API: Iteration
    // =========================================================================

    /// Returns an iterator over all nodes in the database.
    ///
    /// Useful for dump/export operations.
    pub fn iter_nodes(&self) -> impl Iterator<Item = grafeo_core::graph::lpg::Node> + '_ {
        self.store.all_nodes()
    }

    /// Returns an iterator over all edges in the database.
    ///
    /// Useful for dump/export operations.
    pub fn iter_edges(&self) -> impl Iterator<Item = grafeo_core::graph::lpg::Edge> + '_ {
        self.store.all_edges()
    }
}

#[cfg(test)]
mod tests {
    use grafeo_common::types::{EdgeId, NodeId, Value};

    use super::super::GrafeoDB;
    use super::{SNAPSHOT_VERSION, Snapshot, SnapshotEdge, SnapshotNode, SnapshotSchema};

    #[test]
    fn test_restore_snapshot_basic() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();

        // Populate
        session.execute("INSERT (:Person {name: 'Alix'})").unwrap();
        session.execute("INSERT (:Person {name: 'Gus'})").unwrap();

        let snapshot = db.export_snapshot().unwrap();

        // Modify
        session
            .execute("INSERT (:Person {name: 'Vincent'})")
            .unwrap();
        assert_eq!(db.store.node_count(), 3);

        // Restore original
        db.restore_snapshot(&snapshot).unwrap();

        assert_eq!(db.store.node_count(), 2);
        let result = session.execute("MATCH (n:Person) RETURN n.name").unwrap();
        assert_eq!(result.rows.len(), 2);
    }

    #[test]
    fn test_restore_snapshot_validation_failure() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();

        session.execute("INSERT (:Person {name: 'Alix'})").unwrap();

        // Corrupt snapshot: just garbage bytes
        let result = db.restore_snapshot(b"garbage");
        assert!(result.is_err());

        // DB should be unchanged
        assert_eq!(db.store.node_count(), 1);
    }

    #[test]
    fn test_restore_snapshot_empty_db() {
        let db = GrafeoDB::new_in_memory();

        // Export empty snapshot, then populate, then restore to empty
        let empty_snapshot = db.export_snapshot().unwrap();

        let session = db.session();
        session.execute("INSERT (:Person {name: 'Alix'})").unwrap();
        assert_eq!(db.store.node_count(), 1);

        db.restore_snapshot(&empty_snapshot).unwrap();
        assert_eq!(db.store.node_count(), 0);
    }

    #[test]
    fn test_restore_snapshot_with_edges() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();

        session.execute("INSERT (:Person {name: 'Alix'})").unwrap();
        session.execute("INSERT (:Person {name: 'Gus'})").unwrap();
        session
            .execute(
                "MATCH (a:Person {name: 'Alix'}), (b:Person {name: 'Gus'}) INSERT (a)-[:KNOWS]->(b)",
            )
            .unwrap();

        let snapshot = db.export_snapshot().unwrap();
        assert_eq!(db.store.edge_count(), 1);

        // Modify: add more data
        session
            .execute("INSERT (:Person {name: 'Vincent'})")
            .unwrap();

        // Restore
        db.restore_snapshot(&snapshot).unwrap();
        assert_eq!(db.store.node_count(), 2);
        assert_eq!(db.store.edge_count(), 1);
    }

    #[test]
    fn test_restore_snapshot_preserves_sessions() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();

        session.execute("INSERT (:Person {name: 'Alix'})").unwrap();
        let snapshot = db.export_snapshot().unwrap();

        // Modify
        session.execute("INSERT (:Person {name: 'Gus'})").unwrap();

        // Restore
        db.restore_snapshot(&snapshot).unwrap();

        // Session should still work and see restored data
        let result = session.execute("MATCH (n:Person) RETURN n.name").unwrap();
        assert_eq!(result.rows.len(), 1);
    }

    #[test]
    fn test_export_import_roundtrip() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();

        session
            .execute("INSERT (:Person {name: 'Alix', age: 30})")
            .unwrap();

        let snapshot = db.export_snapshot().unwrap();
        let db2 = GrafeoDB::import_snapshot(&snapshot).unwrap();
        let session2 = db2.session();

        let result = session2.execute("MATCH (n:Person) RETURN n.name").unwrap();
        assert_eq!(result.rows.len(), 1);
    }

    // --- to_memory() ---

    #[test]
    fn test_to_memory_empty() {
        let db = GrafeoDB::new_in_memory();
        let copy = db.to_memory().unwrap();
        assert_eq!(copy.store.node_count(), 0);
        assert_eq!(copy.store.edge_count(), 0);
    }

    #[test]
    fn test_to_memory_copies_nodes_and_properties() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();
        session
            .execute("INSERT (:Person {name: 'Alix', age: 30})")
            .unwrap();
        session
            .execute("INSERT (:Person {name: 'Gus', age: 25})")
            .unwrap();

        let copy = db.to_memory().unwrap();
        assert_eq!(copy.store.node_count(), 2);

        let s2 = copy.session();
        let result = s2
            .execute("MATCH (p:Person) RETURN p.name ORDER BY p.name")
            .unwrap();
        assert_eq!(result.rows.len(), 2);
        assert_eq!(result.rows[0][0], Value::String("Alix".into()));
        assert_eq!(result.rows[1][0], Value::String("Gus".into()));
    }

    #[test]
    fn test_to_memory_copies_edges_and_properties() {
        let db = GrafeoDB::new_in_memory();
        let a = db.create_node(&["Person"]);
        db.set_node_property(a, "name", "Alix".into());
        let b = db.create_node(&["Person"]);
        db.set_node_property(b, "name", "Gus".into());
        let edge = db.create_edge(a, b, "KNOWS");
        db.set_edge_property(edge, "since", Value::Int64(2020));

        let copy = db.to_memory().unwrap();
        assert_eq!(copy.store.node_count(), 2);
        assert_eq!(copy.store.edge_count(), 1);

        let s2 = copy.session();
        let result = s2.execute("MATCH ()-[e:KNOWS]->() RETURN e.since").unwrap();
        assert_eq!(result.rows[0][0], Value::Int64(2020));
    }

    #[test]
    fn test_to_memory_is_independent() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();
        session.execute("INSERT (:Person {name: 'Alix'})").unwrap();

        let copy = db.to_memory().unwrap();

        // Mutating original should not affect copy
        session.execute("INSERT (:Person {name: 'Gus'})").unwrap();
        assert_eq!(db.store.node_count(), 2);
        assert_eq!(copy.store.node_count(), 1);
    }

    // --- iter_nodes() / iter_edges() ---

    #[test]
    fn test_iter_nodes_empty() {
        let db = GrafeoDB::new_in_memory();
        assert_eq!(db.iter_nodes().count(), 0);
    }

    #[test]
    fn test_iter_nodes_returns_all() {
        let db = GrafeoDB::new_in_memory();
        let id1 = db.create_node(&["Person"]);
        db.set_node_property(id1, "name", "Alix".into());
        let id2 = db.create_node(&["Animal"]);
        db.set_node_property(id2, "name", "Fido".into());

        let nodes: Vec<_> = db.iter_nodes().collect();
        assert_eq!(nodes.len(), 2);

        let names: Vec<_> = nodes
            .iter()
            .filter_map(|n| n.properties.iter().find(|(k, _)| k.as_str() == "name"))
            .map(|(_, v)| v.clone())
            .collect();
        assert!(names.contains(&Value::String("Alix".into())));
        assert!(names.contains(&Value::String("Fido".into())));
    }

    #[test]
    fn test_iter_edges_empty() {
        let db = GrafeoDB::new_in_memory();
        assert_eq!(db.iter_edges().count(), 0);
    }

    #[test]
    fn test_iter_edges_returns_all() {
        let db = GrafeoDB::new_in_memory();
        let a = db.create_node(&["A"]);
        let b = db.create_node(&["B"]);
        let c = db.create_node(&["C"]);
        db.create_edge(a, b, "R1");
        db.create_edge(b, c, "R2");

        let edges: Vec<_> = db.iter_edges().collect();
        assert_eq!(edges.len(), 2);

        let types: Vec<_> = edges.iter().map(|e| e.edge_type.as_ref()).collect();
        assert!(types.contains(&"R1"));
        assert!(types.contains(&"R2"));
    }

    // --- restore_snapshot() validation ---

    fn make_snapshot(version: u8, nodes: Vec<SnapshotNode>, edges: Vec<SnapshotEdge>) -> Vec<u8> {
        let snap = Snapshot {
            version,
            nodes,
            edges,
            named_graphs: vec![],
            rdf_triples: vec![],
            rdf_named_graphs: vec![],
            schema: SnapshotSchema::default(),
        };
        bincode::serde::encode_to_vec(&snap, bincode::config::standard()).unwrap()
    }

    #[test]
    fn test_restore_rejects_unsupported_version() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();
        session.execute("INSERT (:Person {name: 'Alix'})").unwrap();

        let bytes = make_snapshot(99, vec![], vec![]);

        let result = db.restore_snapshot(&bytes);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("unsupported snapshot version"), "got: {err}");

        // DB unchanged
        assert_eq!(db.store.node_count(), 1);
    }

    #[test]
    fn test_restore_rejects_duplicate_node_ids() {
        let db = GrafeoDB::new_in_memory();
        let session = db.session();
        session.execute("INSERT (:Person {name: 'Alix'})").unwrap();

        let bytes = make_snapshot(
            SNAPSHOT_VERSION,
            vec![
                SnapshotNode {
                    id: NodeId::new(0),
                    labels: vec!["A".into()],
                    properties: vec![],
                },
                SnapshotNode {
                    id: NodeId::new(0),
                    labels: vec!["B".into()],
                    properties: vec![],
                },
            ],
            vec![],
        );

        let result = db.restore_snapshot(&bytes);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("duplicate node ID"), "got: {err}");
        assert_eq!(db.store.node_count(), 1);
    }

    #[test]
    fn test_restore_rejects_duplicate_edge_ids() {
        let db = GrafeoDB::new_in_memory();

        let bytes = make_snapshot(
            SNAPSHOT_VERSION,
            vec![
                SnapshotNode {
                    id: NodeId::new(0),
                    labels: vec![],
                    properties: vec![],
                },
                SnapshotNode {
                    id: NodeId::new(1),
                    labels: vec![],
                    properties: vec![],
                },
            ],
            vec![
                SnapshotEdge {
                    id: EdgeId::new(0),
                    src: NodeId::new(0),
                    dst: NodeId::new(1),
                    edge_type: "REL".into(),
                    properties: vec![],
                },
                SnapshotEdge {
                    id: EdgeId::new(0),
                    src: NodeId::new(0),
                    dst: NodeId::new(1),
                    edge_type: "REL".into(),
                    properties: vec![],
                },
            ],
        );

        let result = db.restore_snapshot(&bytes);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("duplicate edge ID"), "got: {err}");
    }

    #[test]
    fn test_restore_rejects_dangling_source() {
        let db = GrafeoDB::new_in_memory();

        let bytes = make_snapshot(
            SNAPSHOT_VERSION,
            vec![SnapshotNode {
                id: NodeId::new(0),
                labels: vec![],
                properties: vec![],
            }],
            vec![SnapshotEdge {
                id: EdgeId::new(0),
                src: NodeId::new(999),
                dst: NodeId::new(0),
                edge_type: "REL".into(),
                properties: vec![],
            }],
        );

        let result = db.restore_snapshot(&bytes);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("non-existent source node"), "got: {err}");
    }

    #[test]
    fn test_restore_rejects_dangling_destination() {
        let db = GrafeoDB::new_in_memory();

        let bytes = make_snapshot(
            SNAPSHOT_VERSION,
            vec![SnapshotNode {
                id: NodeId::new(0),
                labels: vec![],
                properties: vec![],
            }],
            vec![SnapshotEdge {
                id: EdgeId::new(0),
                src: NodeId::new(0),
                dst: NodeId::new(999),
                edge_type: "REL".into(),
                properties: vec![],
            }],
        );

        let result = db.restore_snapshot(&bytes);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("non-existent destination node"), "got: {err}");
    }
}