use crate::error::Result;
use crate::types::Config;
use crate::{AsyncMemoryGraph, MemoryGraph};
pub struct MigrationHelper;
impl MigrationHelper {
pub async fn verify_compatibility(config: &Config) -> Result<CompatibilityReport> {
let mut report = CompatibilityReport {
sync_accessible: false,
async_accessible: false,
sync_node_count: 0,
async_node_count: 0,
sync_edge_count: 0,
async_edge_count: 0,
compatible: false,
};
match MemoryGraph::open(config.clone()) {
Ok(graph) => {
report.sync_accessible = true;
if let Ok(stats) = graph.stats() {
report.sync_node_count = stats.node_count;
report.sync_edge_count = stats.edge_count;
}
}
Err(_) => return Ok(report),
}
match AsyncMemoryGraph::open(config.clone()).await {
Ok(graph) => {
report.async_accessible = true;
if let Ok(stats) = graph.stats().await {
report.async_node_count = stats.node_count;
report.async_edge_count = stats.edge_count;
}
}
Err(_) => return Ok(report),
}
report.compatible = report.sync_accessible
&& report.async_accessible
&& report.sync_node_count == report.async_node_count
&& report.sync_edge_count == report.async_edge_count;
Ok(report)
}
pub async fn create_checkpoint(config: &Config) -> Result<MigrationCheckpoint> {
let graph = AsyncMemoryGraph::open(config.clone()).await?;
let stats = graph.stats().await?;
Ok(MigrationCheckpoint {
timestamp: chrono::Utc::now(),
node_count: stats.node_count,
edge_count: stats.edge_count,
session_count: stats.session_count,
storage_bytes: stats.storage_bytes,
})
}
pub async fn verify_checkpoint(
config: &Config,
checkpoint: &MigrationCheckpoint,
) -> Result<CheckpointVerification> {
let graph = AsyncMemoryGraph::open(config.clone()).await?;
let stats = graph.stats().await?;
let verification = CheckpointVerification {
valid: stats.node_count >= checkpoint.node_count
&& stats.edge_count >= checkpoint.edge_count,
node_count_match: stats.node_count == checkpoint.node_count,
edge_count_match: stats.edge_count == checkpoint.edge_count,
nodes_added: stats.node_count.saturating_sub(checkpoint.node_count),
edges_added: stats.edge_count.saturating_sub(checkpoint.edge_count),
};
Ok(verification)
}
pub async fn run_migration_test(config: &Config) -> Result<MigrationTestReport> {
let mut report = MigrationTestReport {
success: false,
steps_completed: Vec::new(),
errors: Vec::new(),
};
match Self::create_checkpoint(config).await {
Ok(_) => report
.steps_completed
.push("Checkpoint created".to_string()),
Err(e) => {
report
.errors
.push(format!("Failed to create checkpoint: {}", e));
return Ok(report);
}
}
match MemoryGraph::open(config.clone()) {
Ok(_) => report
.steps_completed
.push("Sync API accessible".to_string()),
Err(e) => {
report.errors.push(format!("Sync API failed: {}", e));
return Ok(report);
}
}
match AsyncMemoryGraph::open(config.clone()).await {
Ok(_) => report
.steps_completed
.push("Async API accessible".to_string()),
Err(e) => {
report.errors.push(format!("Async API failed: {}", e));
return Ok(report);
}
}
match Self::verify_compatibility(config).await {
Ok(compat) if compat.compatible => {
report.steps_completed.push("APIs compatible".to_string());
}
Ok(_) => {
report.errors.push("APIs not compatible".to_string());
return Ok(report);
}
Err(e) => {
report
.errors
.push(format!("Compatibility check failed: {}", e));
return Ok(report);
}
}
report.success = true;
Ok(report)
}
}
#[derive(Debug, Clone)]
pub struct CompatibilityReport {
pub sync_accessible: bool,
pub async_accessible: bool,
pub sync_node_count: u64,
pub async_node_count: u64,
pub sync_edge_count: u64,
pub async_edge_count: u64,
pub compatible: bool,
}
#[derive(Debug, Clone)]
pub struct MigrationCheckpoint {
pub timestamp: chrono::DateTime<chrono::Utc>,
pub node_count: u64,
pub edge_count: u64,
pub session_count: u64,
pub storage_bytes: u64,
}
#[derive(Debug, Clone)]
pub struct CheckpointVerification {
pub valid: bool,
pub node_count_match: bool,
pub edge_count_match: bool,
pub nodes_added: u64,
pub edges_added: u64,
}
#[derive(Debug, Clone)]
pub struct MigrationTestReport {
pub success: bool,
pub steps_completed: Vec<String>,
pub errors: Vec<String>,
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[tokio::test]
async fn test_verify_compatibility_empty_db() {
let dir = tempdir().unwrap();
let config = Config::new(dir.path());
let report = MigrationHelper::verify_compatibility(&config)
.await
.unwrap();
assert!(report.sync_accessible);
assert!(report.async_accessible);
assert!(report.compatible);
assert_eq!(report.sync_node_count, 0);
assert_eq!(report.async_node_count, 0);
}
#[tokio::test]
async fn test_verify_compatibility_with_data() {
let dir = tempdir().unwrap();
let config = Config::new(dir.path());
{
let graph = MemoryGraph::open(config.clone()).unwrap();
graph.create_session().unwrap();
}
let report = MigrationHelper::verify_compatibility(&config)
.await
.unwrap();
assert!(report.sync_accessible);
assert!(report.async_accessible);
assert!(report.compatible);
assert_eq!(report.sync_node_count, 1);
assert_eq!(report.async_node_count, 1);
}
#[tokio::test]
async fn test_create_checkpoint() {
let dir = tempdir().unwrap();
let config = Config::new(dir.path());
{
let graph = AsyncMemoryGraph::open(config.clone()).await.unwrap();
graph.create_session().await.unwrap();
}
let checkpoint = MigrationHelper::create_checkpoint(&config).await.unwrap();
assert_eq!(checkpoint.node_count, 1);
assert_eq!(checkpoint.edge_count, 0);
}
#[tokio::test]
async fn test_verify_checkpoint() {
let dir = tempdir().unwrap();
let config = Config::new(dir.path());
{
let graph = AsyncMemoryGraph::open(config.clone()).await.unwrap();
graph.create_session().await.unwrap();
}
let checkpoint = MigrationHelper::create_checkpoint(&config).await.unwrap();
{
let graph = AsyncMemoryGraph::open(config.clone()).await.unwrap();
graph.create_session().await.unwrap();
}
let verification = MigrationHelper::verify_checkpoint(&config, &checkpoint)
.await
.unwrap();
assert!(verification.valid);
assert!(!verification.node_count_match); assert_eq!(verification.nodes_added, 1);
}
#[tokio::test]
async fn test_run_migration_test() {
let dir = tempdir().unwrap();
let config = Config::new(dir.path());
{
let graph = MemoryGraph::open(config.clone()).unwrap();
graph.create_session().unwrap();
}
let report = MigrationHelper::run_migration_test(&config).await.unwrap();
assert!(report.success, "Migration test should succeed");
assert!(report.errors.is_empty(), "Should have no errors");
assert!(
report.steps_completed.len() >= 4,
"Should complete all steps"
);
}
#[tokio::test]
async fn test_sync_async_interop() {
let dir = tempdir().unwrap();
let config = Config::new(dir.path());
let session_id = {
let graph = MemoryGraph::open(config.clone()).unwrap();
let session = graph.create_session().unwrap();
session.id
};
{
let graph = AsyncMemoryGraph::open(config.clone()).await.unwrap();
let session = graph.get_session(session_id).await.unwrap();
assert_eq!(session.id, session_id);
}
let prompt_id = {
let graph = AsyncMemoryGraph::open(config.clone()).await.unwrap();
graph
.add_prompt(session_id, "Test prompt".to_string(), None)
.await
.unwrap()
};
{
let graph = MemoryGraph::open(config.clone()).unwrap();
let node = graph.get_node(prompt_id).unwrap();
assert_eq!(node.id(), prompt_id);
}
}
}