memorygraphMCP 0.11.7__py3-none-any.whl

memorygraph/migration/scripts/multitenancy_migration.py

@@ -0,0 +1,452 @@
+"""
+001_add_multitenancy - Migration to add multi-tenancy support.
+
+This migration:
+1. Backfills tenant_id for existing memories (if specified)
+2. Sets visibility to 'team' for existing memories
+3. Creates multi-tenant indexes if not already present
+4. Supports rollback by removing tenant assignments
+
+Usage:
+    from memorygraph.migration.scripts import migrate_to_multitenant
+
+    # Migrate with default tenant
+    await migrate_to_multitenant(backend, tenant_id="default")
+
+    # Rollback
+    await rollback_from_multitenant(backend)
+"""
+
+import json
+import logging
+import re
+from typing import Optional
+from ...backends.base import GraphBackend
+from ...backends.sqlite_fallback import SQLiteFallbackBackend
+from ...backends.turso import TursoBackend
+from ...models import DatabaseConnectionError, Memory
+
+logger = logging.getLogger(__name__)
+
+
+async def migrate_to_multitenant(
+    backend: GraphBackend,
+    tenant_id: str = "default",
+    dry_run: bool = False,
+    visibility: str = "team"
+) -> dict:
+    """
+    Migrate existing single-tenant database to multi-tenant mode.
+
+    This function backfills tenant_id and visibility fields for all existing
+    memories that don't have them set.
+
+    Args:
+        backend: Backend instance (must be connected)
+        tenant_id: Tenant ID to assign to existing memories (default: "default")
+        dry_run: If True, only report what would be changed without making changes
+        visibility: Visibility level to set (default: "team")
+
+    Returns:
+        Dictionary with migration statistics:
+        {
+            "success": bool,
+            "dry_run": bool,
+            "memories_updated": int,
+            "errors": list
+        }
+
+    Raises:
+        DatabaseConnectionError: If backend is not connected
+        ValueError: If tenant_id is empty or visibility is invalid
+
+    Example:
+        >>> backend = SQLiteFallbackBackend()
+        >>> await backend.connect()
+        >>> result = await migrate_to_multitenant(backend, tenant_id="acme-corp")
+        >>> print(f"Updated {result['memories_updated']} memories")
+    """
+    # Check backend connection (use _connected attribute for SQLite backends)
+    is_connected = getattr(backend, '_connected', False)
+    if not backend or not is_connected:
+        raise DatabaseConnectionError("Backend must be connected before migration")
+
+    # Validate tenant_id is not empty
+    if not tenant_id or not tenant_id.strip():
+        raise ValueError("tenant_id cannot be empty")
+
+    # Validate tenant_id format (alphanumeric, dashes, underscores only)
+    if not re.match(r'^[a-zA-Z0-9-_]+$', tenant_id):
+        raise ValueError(
+            "tenant_id must contain only alphanumeric characters, dashes, and underscores"
+        )
+
+    # Validate tenant_id length
+    if len(tenant_id) > 64:
+        raise ValueError("tenant_id must be 64 characters or less")
+
+    # Validate visibility value
+    valid_visibility = ["private", "project", "team", "public"]
+    if visibility not in valid_visibility:
+        raise ValueError(f"visibility must be one of {valid_visibility}, got '{visibility}'")
+
+    logger.info(f"Starting multi-tenancy migration (dry_run={dry_run})")
+    logger.info(f"Assigning tenant_id='{tenant_id}', visibility='{visibility}'")
+
+    errors = []
+    memories_updated = 0
+
+    try:
+        # SQLite-based backends (SQLite, Turso) - use duck typing to avoid
+        # isinstance issues when modules are reloaded during testing.
+        # SQLite backends have a 'conn' attribute for the database connection.
+        if hasattr(backend, 'conn') and backend.conn is not None:
+            memories_updated = await _migrate_sqlite_backend(
+                backend, tenant_id, visibility, dry_run
+            )
+
+        # Neo4j/Memgraph backends - use execute_query method
+        elif hasattr(backend, 'execute_query'):
+            memories_updated = await _migrate_graph_backend(
+                backend, tenant_id, visibility, dry_run
+            )
+
+        else:
+            raise ValueError(f"Unsupported backend type: {type(backend).__name__}")
+
+        logger.info(f"Migration completed: {memories_updated} memories processed")
+
+        return {
+            "success": True,
+            "dry_run": dry_run,
+            "memories_updated": memories_updated,
+            "tenant_id": tenant_id,
+            "visibility": visibility,
+            "errors": errors
+        }
+
+    except Exception as e:
+        logger.error(f"Migration failed: {e}")
+        errors.append(str(e))
+        return {
+            "success": False,
+            "dry_run": dry_run,
+            "memories_updated": memories_updated,
+            "tenant_id": tenant_id,
+            "visibility": visibility,
+            "errors": errors
+        }
+
+
+async def _migrate_sqlite_backend(
+    backend: SQLiteFallbackBackend,
+    tenant_id: str,
+    visibility: str,
+    dry_run: bool
+) -> int:
+    """
+    Migrate SQLite-based backend to multi-tenant mode.
+
+    The properties are stored in a flat structure with context fields prefixed
+    with 'context_' (e.g., 'context_tenant_id', 'context_visibility').
+
+    Args:
+        backend: SQLite backend instance
+        tenant_id: Tenant ID to assign
+        visibility: Visibility level to set
+        dry_run: If True, only count without updating
+
+    Returns:
+        Number of memories updated
+    """
+    cursor = backend.conn.cursor()
+
+    # Count memories without tenant_id (using flat property structure)
+    # Properties use 'context_tenant_id' not 'context.tenant_id'
+    cursor.execute("""
+        SELECT COUNT(*) FROM nodes
+        WHERE label = 'Memory'
+        AND (
+            json_extract(properties, '$.context_tenant_id') IS NULL
+            OR json_extract(properties, '$.context_tenant_id') = ''
+        )
+    """)
+
+    count = cursor.fetchone()[0]
+    logger.info(f"Found {count} memories without tenant_id")
+
+    if dry_run:
+        logger.info(f"DRY RUN: Would update {count} memories")
+        return count
+
+    # Update memories by fetching, modifying, and updating each one
+    # This ensures proper JSON structure handling
+    cursor.execute("""
+        SELECT id, properties FROM nodes
+        WHERE label = 'Memory'
+        AND (
+            json_extract(properties, '$.context_tenant_id') IS NULL
+            OR json_extract(properties, '$.context_tenant_id') = ''
+        )
+    """)
+
+    updated = 0
+
+    for row in cursor.fetchall():
+        node_id = row[0]
+        properties = json.loads(row[1])
+
+        # Set tenant_id and visibility using flat property structure
+        properties['context_tenant_id'] = tenant_id
+        properties['context_visibility'] = visibility
+
+        # Update the node
+        cursor.execute("""
+            UPDATE nodes
+            SET properties = ?,
+                updated_at = CURRENT_TIMESTAMP
+            WHERE id = ?
+        """, [json.dumps(properties), node_id])
+
+        updated += 1
+
+    backend.conn.commit()
+    logger.info(f"Updated {updated} memories with tenant_id='{tenant_id}'")
+
+    return updated
+
+
+async def _migrate_graph_backend(
+    backend: GraphBackend,
+    tenant_id: str,
+    visibility: str,
+    dry_run: bool
+) -> int:
+    """
+    Migrate graph-based backend (Neo4j/Memgraph) to multi-tenant mode.
+
+    Args:
+        backend: Graph backend instance
+        tenant_id: Tenant ID to assign
+        visibility: Visibility level to set
+        dry_run: If True, only count without updating
+
+    Returns:
+        Number of memories updated
+    """
+    # Count memories without tenant_id
+    count_query = """
+        MATCH (m:Memory)
+        WHERE m.context_tenant_id IS NULL OR m.context_tenant_id = ''
+        RETURN count(m) as count
+    """
+
+    count_result = await backend.execute_query(count_query)
+    count = count_result[0]['count'] if count_result else 0
+
+    logger.info(f"Found {count} memories without tenant_id")
+
+    if dry_run:
+        logger.info(f"DRY RUN: Would update {count} memories")
+        return count
+
+    # Update memories with tenant_id and visibility
+    update_query = """
+        MATCH (m:Memory)
+        WHERE m.context_tenant_id IS NULL OR m.context_tenant_id = ''
+        SET m.context_tenant_id = $tenant_id,
+            m.context_visibility = $visibility,
+            m.updated_at = timestamp()
+        RETURN count(m) as updated
+    """
+
+    result = await backend.execute_query(
+        update_query,
+        {"tenant_id": tenant_id, "visibility": visibility},
+        write=True
+    )
+
+    updated = result[0]['updated'] if result else 0
+    logger.info(f"Updated {updated} memories with tenant_id='{tenant_id}'")
+
+    return updated
+
+
+async def rollback_from_multitenant(
+    backend: GraphBackend,
+    dry_run: bool = False
+) -> dict:
+    """
+    Rollback multi-tenancy migration by removing tenant_id assignments.
+
+    NOTE: This does not delete the tenant_id fields, it only sets them to NULL.
+    This preserves the option to re-enable multi-tenancy in the future.
+
+    Args:
+        backend: Backend instance (must be connected)
+        dry_run: If True, only report what would be changed
+
+    Returns:
+        Dictionary with rollback statistics
+
+    Example:
+        >>> result = await rollback_from_multitenant(backend)
+        >>> print(f"Rolled back {result['memories_updated']} memories")
+    """
+    # Check backend connection (use _connected attribute for SQLite backends)
+    is_connected = getattr(backend, '_connected', False)
+    if not backend or not is_connected:
+        raise DatabaseConnectionError("Backend must be connected before rollback")
+
+    logger.info(f"Starting multi-tenancy rollback (dry_run={dry_run})")
+
+    errors = []
+    memories_updated = 0
+
+    try:
+        # SQLite-based backends - use duck typing (check for conn attribute)
+        if hasattr(backend, 'conn') and backend.conn is not None:
+            memories_updated = await _rollback_sqlite_backend(backend, dry_run)
+
+        # Graph backends - use execute_query method
+        elif hasattr(backend, 'execute_query'):
+            memories_updated = await _rollback_graph_backend(backend, dry_run)
+
+        else:
+            raise ValueError(f"Unsupported backend type: {type(backend).__name__}")
+
+        logger.info(f"Rollback completed: {memories_updated} memories processed")
+
+        return {
+            "success": True,
+            "dry_run": dry_run,
+            "memories_updated": memories_updated,
+            "errors": errors
+        }
+
+    except Exception as e:
+        logger.error(f"Rollback failed: {e}")
+        errors.append(str(e))
+        return {
+            "success": False,
+            "dry_run": dry_run,
+            "memories_updated": memories_updated,
+            "errors": errors
+        }
+
+
+async def _rollback_sqlite_backend(
+    backend: SQLiteFallbackBackend,
+    dry_run: bool
+) -> int:
+    """
+    Rollback SQLite backend from multi-tenant mode.
+
+    The properties are stored in a flat structure with context fields prefixed
+    with 'context_' (e.g., 'context_tenant_id', 'context_visibility').
+
+    Args:
+        backend: SQLite backend instance
+        dry_run: If True, only count without updating
+
+    Returns:
+        Number of memories updated
+    """
+    cursor = backend.conn.cursor()
+
+    # Count memories with tenant_id (using flat property structure)
+    cursor.execute("""
+        SELECT COUNT(*) FROM nodes
+        WHERE label = 'Memory'
+        AND json_extract(properties, '$.context_tenant_id') IS NOT NULL
+        AND json_extract(properties, '$.context_tenant_id') != ''
+    """)
+
+    count = cursor.fetchone()[0]
+    logger.info(f"Found {count} memories with tenant_id")
+
+    if dry_run:
+        logger.info(f"DRY RUN: Would clear tenant_id from {count} memories")
+        return count
+
+    # Clear tenant_id from memories
+    cursor.execute("""
+        SELECT id, properties FROM nodes
+        WHERE label = 'Memory'
+        AND json_extract(properties, '$.context_tenant_id') IS NOT NULL
+        AND json_extract(properties, '$.context_tenant_id') != ''
+    """)
+
+    updated = 0
+
+    for row in cursor.fetchall():
+        node_id = row[0]
+        properties = json.loads(row[1])
+
+        # Clear tenant_id (set to NULL) using flat property structure
+        properties['context_tenant_id'] = None
+        # Reset visibility to default
+        properties['context_visibility'] = 'project'
+
+        # Update the node
+        cursor.execute("""
+            UPDATE nodes
+            SET properties = ?,
+                updated_at = CURRENT_TIMESTAMP
+            WHERE id = ?
+        """, [json.dumps(properties), node_id])
+
+        updated += 1
+
+    backend.conn.commit()
+    logger.info(f"Cleared tenant_id from {updated} memories")
+
+    return updated
+
+
+async def _rollback_graph_backend(
+    backend: GraphBackend,
+    dry_run: bool
+) -> int:
+    """
+    Rollback graph backend from multi-tenant mode.
+
+    Args:
+        backend: Graph backend instance
+        dry_run: If True, only count without updating
+
+    Returns:
+        Number of memories updated
+    """
+    # Count memories with tenant_id
+    count_query = """
+        MATCH (m:Memory)
+        WHERE m.context_tenant_id IS NOT NULL
+        RETURN count(m) as count
+    """
+
+    count_result = await backend.execute_query(count_query)
+    count = count_result[0]['count'] if count_result else 0
+
+    logger.info(f"Found {count} memories with tenant_id")
+
+    if dry_run:
+        logger.info(f"DRY RUN: Would clear tenant_id from {count} memories")
+        return count
+
+    # Clear tenant_id from memories
+    update_query = """
+        MATCH (m:Memory)
+        WHERE m.context_tenant_id IS NOT NULL
+        SET m.context_tenant_id = NULL,
+            m.context_visibility = 'project',
+            m.updated_at = timestamp()
+        RETURN count(m) as updated
+    """
+
+    result = await backend.execute_query(update_query, write=True)
+    updated = result[0]['updated'] if result else 0
+
+    logger.info(f"Cleared tenant_id from {updated} memories")
+
+    return updated

memorygraph/migration_tools_module.py

@@ -0,0 +1,146 @@
+"""
+Migration tools module - MCP tool definitions and handlers for database migration.
+"""
+
+from mcp.types import Tool
+from .tools.migration_tools import handle_migrate_database, handle_validate_migration
+
+# Tool definitions for MCP
+MIGRATION_TOOLS = [
+    Tool(
+        name="migrate_database",
+        description="""Migrate memories from current backend to another backend (e.g., SQLite → FalkorDB).
+
+WHEN TO USE:
+- Moving from development (SQLite) to production (FalkorDB, Neo4j)
+- Switching backend providers
+- Disaster recovery to different backend
+- Testing performance across backends
+- Backend consolidation or splitting
+
+HOW TO USE:
+- Always use dry_run=True first to validate
+- Specify target_backend type (sqlite, neo4j, memgraph, falkordb, falkordblite)
+- Provide target_config with connection details
+- Set verify=True to ensure data integrity
+- Migration includes memories and relationships
+
+SAFETY FEATURES:
+- Dry-run mode validates without changes
+- Verification checks data integrity
+- Automatic rollback on failure
+- Progress reporting for large migrations
+
+EXAMPLES:
+- Validate: migrate_database(target_backend="falkordb", target_config={"uri": "redis://prod:6379"}, dry_run=True)
+- Migrate: migrate_database(target_backend="falkordb", target_config={"uri": "redis://prod:6379"}, verify=True)
+- Test: migrate_database(target_backend="sqlite", target_config={"path": "/tmp/test.db"})
+
+RETURNS:
+- success: Boolean indicating if migration succeeded
+- imported_memories: Number of memories migrated
+- imported_relationships: Number of relationships migrated
+- verification: Data integrity check results
+- errors: Any errors encountered""",
+        inputSchema={
+            "type": "object",
+            "properties": {
+                "target_backend": {
+                    "type": "string",
+                    "enum": ["sqlite", "neo4j", "memgraph", "falkordb", "falkordblite"],
+                    "description": "Target backend type to migrate to"
+                },
+                "target_config": {
+                    "type": "object",
+                    "properties": {
+                        "path": {
+                            "type": "string",
+                            "description": "Database path (for sqlite/falkordblite)"
+                        },
+                        "uri": {
+                            "type": "string",
+                            "description": "Database URI (for neo4j/memgraph/falkordb)"
+                        },
+                        "username": {
+                            "type": "string",
+                            "description": "Database username (optional)"
+                        },
+                        "password": {
+                            "type": "string",
+                            "description": "Database password (optional)"
+                        },
+                        "database": {
+                            "type": "string",
+                            "description": "Database name (optional)"
+                        }
+                    },
+                    "description": "Target backend configuration"
+                },
+                "dry_run": {
+                    "type": "boolean",
+                    "default": False,
+                    "description": "Validate without making changes (RECOMMENDED: use true first)"
+                },
+                "skip_duplicates": {
+                    "type": "boolean",
+                    "default": True,
+                    "description": "Skip memories that already exist in target"
+                },
+                "verify": {
+                    "type": "boolean",
+                    "default": True,
+                    "description": "Verify data integrity after migration"
+                }
+            },
+            "required": ["target_backend"]
+        }
+    ),
+    Tool(
+        name="validate_migration",
+        description="""Validate that migration to target backend would succeed without making changes.
+
+This is a convenience wrapper for migrate_database with dry_run=True.
+
+WHEN TO USE:
+- Before running actual migration
+- Checking if target backend is accessible
+- Estimating migration size and duration
+- Validating target configuration
+
+CHECKS PERFORMED:
+- Source backend accessible
+- Target backend accessible
+- Backend compatibility
+- Configuration validity
+- Data export feasibility
+
+EXAMPLES:
+- validate_migration(target_backend="falkordb", target_config={"uri": "redis://prod:6379"})
+- validate_migration(target_backend="neo4j", target_config={"uri": "bolt://localhost:7687", "username": "neo4j", "password": "password"})
+
+RETURNS:
+- Same as migrate_database but with dry_run=True
+- No data is written to target""",
+        inputSchema={
+            "type": "object",
+            "properties": {
+                "target_backend": {
+                    "type": "string",
+                    "enum": ["sqlite", "neo4j", "memgraph", "falkordb", "falkordblite"],
+                    "description": "Target backend type to validate migration to"
+                },
+                "target_config": {
+                    "type": "object",
+                    "description": "Target backend configuration"
+                }
+            },
+            "required": ["target_backend"]
+        }
+    )
+]
+
+# Tool handlers mapping
+MIGRATION_TOOL_HANDLERS = {
+    "migrate_database": handle_migrate_database,
+    "validate_migration": handle_validate_migration
+}