LightRAG/lightrag/kg/redis_tenant_support.py

# Redis Multi-Tenant Support Module

from typing import Any, Dict, List, Optional


class RedisTenantHelper:
    """Helper class for Redis multi-tenant key management"""
    
    SEPARATOR = ":"
    
    @staticmethod
    def make_tenant_key(tenant_id: str, kb_id: str, original_key: str) -> str:
        """
        Create a tenant-scoped Redis key.
        Format: tenant_id:kb_id:original_key
        """
        return f"{tenant_id}{RedisTenantHelper.SEPARATOR}{kb_id}{RedisTenantHelper.SEPARATOR}{original_key}"
    
    @staticmethod
    def parse_tenant_key(tenant_key: str) -> Dict[str, str]:
        """
        Parse a tenant-scoped key to extract components.
        
        Returns:
            Dict with tenant_id, kb_id, and original_key
        """
        parts = tenant_key.split(RedisTenantHelper.SEPARATOR, 2)
        if len(parts) == 3:
            return {
                "tenant_id": parts[0],
                "kb_id": parts[1],
                "original_key": parts[2]
            }
        return {"original_key": tenant_key}
    
    @staticmethod
    def get_tenant_key_pattern(tenant_id: str, kb_id: str, pattern: str = "*") -> str:
        """
        Get a key pattern for scanning tenant-specific keys.
        Useful for operations like SCAN, KEYS.
        
        Format: tenant_id:kb_id:*
        """
        return f"{tenant_id}{RedisTenantHelper.SEPARATOR}{kb_id}{RedisTenantHelper.SEPARATOR}{pattern}"
    
    @staticmethod
    def extract_original_key(tenant_key: str) -> str:
        """Extract the original key from a tenant-scoped key"""
        parts = tenant_key.split(RedisTenantHelper.SEPARATOR, 2)
        return parts[2] if len(parts) == 3 else tenant_key
    
    @staticmethod
    def batch_make_tenant_keys(
        tenant_id: str,
        kb_id: str,
        keys: List[str]
    ) -> List[str]:
        """Create multiple tenant-scoped keys at once"""
        return [
            RedisTenantHelper.make_tenant_key(tenant_id, kb_id, key)
            for key in keys
        ]
    
    @staticmethod
    def batch_parse_tenant_keys(tenant_keys: List[str]) -> List[Dict[str, str]]:
        return [RedisTenantHelper.parse_tenant_key(key) for key in tenant_keys]


class RedisTenantNamespace:
    """Context manager for tenant-scoped Redis operations"""
    
    def __init__(self, redis_client, tenant_id: str, kb_id: str):
        self.redis = redis_client
        self.tenant_id = tenant_id
        self.kb_id = kb_id
        self._prefix = RedisTenantHelper.get_tenant_key_pattern(tenant_id, kb_id, "")
    
    async def get(self, key: str) -> Optional[Any]:
        """Get a value with tenant scoping"""
        tenant_key = RedisTenantHelper.make_tenant_key(self.tenant_id, self.kb_id, key)
        return await self.redis.get(tenant_key)
    
    async def set(self, key: str, value: Any, ex: Optional[int] = None) -> bool:
        """Set a value with tenant scoping"""
        tenant_key = RedisTenantHelper.make_tenant_key(self.tenant_id, self.kb_id, key)
        if ex:
            return await self.redis.setex(tenant_key, ex, value)
        return await self.redis.set(tenant_key, value)
    
    async def delete(self, keys: List[str]) -> int:
        """Delete values with tenant scoping"""
        tenant_keys = RedisTenantHelper.batch_make_tenant_keys(
            self.tenant_id, self.kb_id, keys
        )
        return await self.redis.delete(*tenant_keys)
    
    async def exists(self, key: str) -> bool:
        """Check if key exists with tenant scoping"""
        tenant_key = RedisTenantHelper.make_tenant_key(self.tenant_id, self.kb_id, key)
        return await self.redis.exists(tenant_key) > 0
    
    async def scan_keys(self, pattern: str = "*") -> List[str]:
        """Scan keys for this tenant"""
        tenant_pattern = self._prefix + pattern
        keys = []
        cursor = 0
        while True:
            cursor, batch = await self.redis.scan(cursor, match=tenant_pattern, count=100)
            keys.extend([
                RedisTenantHelper.extract_original_key(k.decode() if isinstance(k, bytes) else k)
                for k in batch
            ])
            if cursor == 0:
                break
        return keys
    
    async def delete_all(self) -> int:
        """Delete all keys for this tenant"""
        keys = await self.scan_keys()
        if keys:
            return await self.delete(keys)
        return 0


# ============================================================================
# MIGRATION HELPER FOR EXISTING REDIS DATA
# ============================================================================

async def migrate_redis_to_tenant(
    redis_client,
    old_key_pattern: str = "*",
    default_tenant_id: str = "default",
    default_kb_id: str = "default",
    dry_run: bool = True
) -> Dict[str, int]:
    """
    Migrate existing Redis keys to tenant-scoped format.
    
    This operation:
    1. Scans all keys matching the pattern
    2. For each key, creates a new tenant-scoped key
    3. Optionally deletes the old key (if not dry_run)
    
    Args:
        redis_client: Redis async client
        old_key_pattern: Pattern to match old keys (default: all keys)
        default_tenant_id: Tenant ID to use for migration
        default_kb_id: KB ID to use for migration
        dry_run: If True, don't actually delete old keys
    
    Returns:
        Migration statistics
    """
    stats = {
        "migrated": 0,
        "failed": 0,
        "skipped": 0
    }
    
    cursor = 0
    while True:
        cursor, keys = await redis_client.scan(cursor, match=old_key_pattern, count=100)
        
        for key_bytes in keys:
            key = key_bytes.decode() if isinstance(key_bytes, bytes) else key_bytes
            
            # Check if key is already tenant-scoped (skip if yes)
            parts = key.split(":", 2)
            if len(parts) == 3:
                stats["skipped"] += 1
                continue
            
            try:
                # Get the value
                value = await redis_client.get(key)
                if value is None:
                    stats["skipped"] += 1
                    continue
                
                # Create tenant-scoped key
                tenant_key = RedisTenantHelper.make_tenant_key(
                    default_tenant_id,
                    default_kb_id,
                    key
                )
                
                # Set the new key
                await redis_client.set(tenant_key, value)
                
                # Delete old key (unless dry_run)
                if not dry_run:
                    await redis_client.delete(key)
                
                stats["migrated"] += 1
            
            except Exception as e:
                print(f"Error migrating key {key}: {e}")
                stats["failed"] += 1
        
        if cursor == 0:
            break
    
    return stats