contextsynapse 1.0.0__py3-none-any.whl

contextsynapse/storage/lmdb_graph_storage.py

@@ -0,0 +1,366 @@
+"""
+LMDB-backed graph storage with CSR-compatible API.
+
+Uses msgpack for fast serialization. Supports node CRUD with an in-memory
+type index rebuilt on startup. Edges are stored but not yet queried by this
+module — the edge sub-database is reserved for future use.
+
+Environment variable:
+    AICONTEXTDB_LMDB_MAP_SIZE — map size in bytes (default: 2 GB)
+"""
+from __future__ import annotations
+
+import copy
+import os
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterator, List, Optional
+
+try:
+    import lmdb
+    import msgpack
+    LMDB_AVAILABLE = True
+except ImportError:
+    LMDB_AVAILABLE = False
+
+# 2 GB default map size
+_DEFAULT_MAP_SIZE = int(os.environ.get("CONTEXTSYNAPSE_LMDB_MAP_SIZE") or os.environ.get("AICONTEXTDB_LMDB_MAP_SIZE", 2 * 1024 ** 3))
+
+# Sub-database names (bytes)
+_DB_NODES = b"nodes"
+_DB_EDGES = b"edges"
+_DB_EDGE_IDX = b"edge_idx"
+_DB_META = b"meta"
+
+
+@dataclass
+class LMDBNode:
+    id: str
+    node_type: str
+    properties: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def label(self) -> str:
+        return self.node_type
+
+
+class LMDBGraphStorage:
+    """
+    Persistent graph node store backed by LMDB.
+
+    Sub-databases:
+        nodes    — node_id (bytes) → msgpack({t, p})
+        edges    — edge_id (bytes) → msgpack edge record (future)
+        edge_idx — composite index for adjacency lookups (future)
+        meta     — store metadata (future)
+    """
+
+    def __init__(
+        self,
+        path: str,
+        map_size: Optional[int] = None,
+        namespace: str = "",
+    ) -> None:
+        if not LMDB_AVAILABLE:
+            raise RuntimeError(
+                "lmdb and msgpack are required for LMDBGraphStorage. "
+                "Install them with: pip install lmdb msgpack"
+            )
+
+        self._path = path
+        self._namespace = namespace
+        self._map_size = map_size if map_size is not None else _DEFAULT_MAP_SIZE
+
+        os.makedirs(path, exist_ok=True)
+
+        self._env = lmdb.open(
+            path,
+            map_size=self._map_size,
+            max_dbs=4,
+            writemap=True,
+            metasync=False,
+            sync=False,
+        )
+
+        # Open sub-databases
+        with self._env.begin(write=True) as txn:
+            self._db_nodes = self._env.open_db(_DB_NODES, txn=txn)
+            self._db_edges = self._env.open_db(_DB_EDGES, txn=txn)
+            self._db_edge_idx = self._env.open_db(_DB_EDGE_IDX, txn=txn)
+            self._db_meta = self._env.open_db(_DB_META, txn=txn)
+
+        # In-memory type index: type_str -> set of node_ids
+        self._type_index: Dict[str, set] = {}
+        self._build_type_index()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _encode_key(self, node_id: str) -> bytes:
+        return node_id.encode("utf-8")
+
+    def _pack_node(self, node_type: str, properties: Dict[str, Any]) -> bytes:
+        return msgpack.packb({"t": node_type, "p": properties}, use_bin_type=True)
+
+    def _unpack_node(self, node_id: str, raw: bytes) -> LMDBNode:
+        data = msgpack.unpackb(raw, raw=False)
+        return LMDBNode(
+            id=node_id,
+            node_type=data["t"],
+            properties=dict(data.get("p") or {}),
+        )
+
+    def _build_type_index(self) -> None:
+        """Scan all nodes on startup and populate the in-memory type index."""
+        self._type_index = {}
+        with self._env.begin(db=self._db_nodes) as txn:
+            cursor = txn.cursor()
+            for raw_key, raw_val in cursor.iternext_dup() if False else cursor:
+                node_id = raw_key.decode("utf-8")
+                data = msgpack.unpackb(raw_val, raw=False)
+                ntype = data["t"]
+                self._type_index.setdefault(ntype, set()).add(node_id)
+
+    def _index_add(self, node_type: str, node_id: str) -> None:
+        self._type_index.setdefault(node_type, set()).add(node_id)
+
+    def _index_remove(self, node_type: str, node_id: str) -> None:
+        if node_type in self._type_index:
+            self._type_index[node_type].discard(node_id)
+            if not self._type_index[node_type]:
+                del self._type_index[node_type]
+
+    def _check_resize(self) -> None:
+        """Auto-resize the map if usage exceeds 80% of the current map size."""
+        info = self._env.info()
+        stat = self._env.stat()
+        # info() exposes last_pgno and map_size; stat() exposes psize (page size)
+        page_size = stat.get("psize", 4096)
+        used = info["last_pgno"] * page_size
+        capacity = info.get("map_size", self._map_size)
+        if capacity > 0 and used / capacity > 0.8:
+            new_size = capacity * 2
+            self._map_size = new_size
+            self._env.set_mapsize(new_size)
+
+    # ------------------------------------------------------------------
+    # Node CRUD
+    # ------------------------------------------------------------------
+
+    def add_node(
+        self,
+        node_id: str,
+        node_type: str,
+        properties: Optional[Dict[str, Any]] = None,
+        alias: Optional[str] = None,
+    ) -> bool:
+        """
+        Add a node. Returns False (without modifying storage) if the ID
+        already exists. Returns True on success.
+        """
+        self._check_resize()
+        key = self._encode_key(node_id)
+        props = properties or {}
+
+        with self._env.begin(write=True, db=self._db_nodes) as txn:
+            existing = txn.get(key)
+            if existing is not None:
+                return False
+            txn.put(key, self._pack_node(node_type, props))
+
+        self._index_add(node_type, node_id)
+        return True
+
+    def get_node(self, node_id: str) -> Optional[LMDBNode]:
+        """
+        Retrieve a node by ID. Returns a copy (mutations do not affect storage).
+        Returns None if not found.
+        """
+        key = self._encode_key(node_id)
+        with self._env.begin(db=self._db_nodes) as txn:
+            raw = txn.get(key)
+            if raw is None:
+                return None
+            node = self._unpack_node(node_id, bytes(raw))  # copy out of mmap
+        # Return a deep copy so callers cannot mutate storage state via the object
+        return copy.deepcopy(node)
+
+    def update_node_properties(self, node_id: str, updates: Dict[str, Any]) -> None:
+        """
+        Merge *updates* into the existing node's properties (read-modify-write
+        in a single transaction). No-op if the node does not exist.
+        """
+        key = self._encode_key(node_id)
+        with self._env.begin(write=True, db=self._db_nodes) as txn:
+            raw = txn.get(key)
+            if raw is None:
+                return
+            data = msgpack.unpackb(bytes(raw), raw=False)
+            props = dict(data.get("p") or {})
+            props.update(updates)
+            txn.put(key, self._pack_node(data["t"], props))
+
+    def delete_node(self, node_id: str) -> bool:
+        """
+        Delete a node and its connected edges.
+        Returns True if the node existed and was deleted, False otherwise.
+        """
+        key = self._encode_key(node_id)
+        with self._env.begin(write=True, db=self._db_nodes) as txn:
+            raw = txn.get(key)
+            if raw is None:
+                return False
+            data = msgpack.unpackb(bytes(raw), raw=False)
+            node_type = data["t"]
+            txn.delete(key)
+
+        self._index_remove(node_type, node_id)
+        self._delete_edges_for_node(node_id)
+        return True
+
+    def _delete_edges_for_node(self, node_id: str) -> None:
+        """Delete all edges where source or target matches *node_id*."""
+        prefix_as_source = (node_id + "\x00").encode("utf-8")
+        to_delete: List[bytes] = []
+
+        with self._env.begin(db=self._db_edges) as txn:
+            cursor = txn.cursor()
+            if cursor.first():
+                while True:
+                    raw_key = bytes(cursor.key())
+                    parts = raw_key.split(b"\x00")
+                    # parts: [source, target, label]
+                    if len(parts) == 3:
+                        src = parts[0].decode("utf-8")
+                        tgt = parts[1].decode("utf-8")
+                        if src == node_id or tgt == node_id:
+                            to_delete.append(raw_key)
+                    if not cursor.next():
+                        break
+
+        if to_delete:
+            with self._env.begin(write=True, db=self._db_edges) as txn:
+                for k in to_delete:
+                    txn.delete(k)
+
+    # ------------------------------------------------------------------
+    # Edge CRUD
+    # ------------------------------------------------------------------
+
+    def _encode_edge_key(self, source_id: str, target_id: str, label: str) -> bytes:
+        """Composite edge key: source\x00target\x00label (all utf-8)."""
+        return "\x00".join([source_id, target_id, label]).encode("utf-8")
+
+    def add_edge(
+        self,
+        source_id: str,
+        target_id: str,
+        label: str,
+        properties: Optional[Dict[str, Any]] = None,
+    ) -> bool:
+        """
+        Add a directed edge from *source_id* to *target_id* with *label*.
+        Both nodes must already exist; returns False otherwise.
+        Returns True on success.
+        """
+        # Verify both nodes exist
+        src_key = self._encode_key(source_id)
+        tgt_key = self._encode_key(target_id)
+        with self._env.begin(db=self._db_nodes) as txn:
+            if txn.get(src_key) is None or txn.get(tgt_key) is None:
+                return False
+
+        edge_key = self._encode_edge_key(source_id, target_id, label)
+        packed = msgpack.packb({"p": properties or {}}, use_bin_type=True)
+        with self._env.begin(write=True, db=self._db_edges) as txn:
+            txn.put(edge_key, packed)
+        return True
+
+    def get_edges(self, node_id: Optional[str] = None) -> List[Dict[str, Any]]:
+        """
+        Return a list of edge dicts.  Each dict has keys:
+            source, target, label, properties
+        If *node_id* is given, only edges where source OR target equals
+        *node_id* are returned.  If None, all edges are returned.
+        """
+        results: List[Dict[str, Any]] = []
+        with self._env.begin(db=self._db_edges) as txn:
+            cursor = txn.cursor()
+            if not cursor.first():
+                return results
+            while True:
+                raw_key = bytes(cursor.key())
+                raw_val = bytes(cursor.value())
+                parts = raw_key.split(b"\x00", 2)
+                if len(parts) == 3:
+                    src = parts[0].decode("utf-8")
+                    tgt = parts[1].decode("utf-8")
+                    lbl = parts[2].decode("utf-8")
+                    if node_id is None or src == node_id or tgt == node_id:
+                        data = msgpack.unpackb(raw_val, raw=False)
+                        results.append({
+                            "source": src,
+                            "target": tgt,
+                            "label": lbl,
+                            "properties": dict(data.get("p") or {}),
+                        })
+                if not cursor.next():
+                    break
+        return results
+
+    def get_edge_count(self) -> int:
+        """Return the number of edges via LMDB stat (O(1))."""
+        with self._env.begin(db=self._db_edges) as txn:
+            return txn.stat()["entries"]
+
+    def remove_edge(self, source_id: str, target_id: str, label: str) -> bool:
+        """
+        Delete the edge identified by the composite key.
+        Returns True if deleted, False if not found.
+        """
+        edge_key = self._encode_edge_key(source_id, target_id, label)
+        with self._env.begin(write=True, db=self._db_edges) as txn:
+            existing = txn.get(edge_key)
+            if existing is None:
+                return False
+            txn.delete(edge_key)
+        return True
+
+    def get_all_nodes(self) -> Iterator[LMDBNode]:
+        """
+        Lazy cursor scan over all nodes. Yields one LMDBNode at a time.
+        Each yielded node is a copy (safe for mutation by callers).
+        """
+        with self._env.begin(db=self._db_nodes) as txn:
+            cursor = txn.cursor()
+            for raw_key, raw_val in cursor:
+                node_id = raw_key.decode("utf-8")
+                node = self._unpack_node(node_id, bytes(raw_val))
+                yield copy.deepcopy(node)
+
+    def get_node_count(self) -> int:
+        """Return the number of nodes via LMDB stat (O(1))."""
+        with self._env.begin(db=self._db_nodes) as txn:
+            return txn.stat()["entries"]
+
+    # ------------------------------------------------------------------
+    # Stats & lifecycle
+    # ------------------------------------------------------------------
+
+    def stats(self) -> dict:
+        """Return a summary dict about the store."""
+        info = self._env.info()
+        return {
+            "namespace": self._namespace,
+            "path": self._path,
+            "node_count": self.get_node_count(),
+            "map_size": self._map_size,
+            "last_pgno": info.get("last_pgno"),
+            "page_size": info.get("page_size"),
+            "type_index": {k: len(v) for k, v in self._type_index.items()},
+        }
+
+    def close(self) -> None:
+        """Flush and close the LMDB environment."""
+        self._env.sync()
+        self._env.close()

contextsynapse/storage/migration.py

@@ -0,0 +1,279 @@
+"""
+Migration System for AIContextDB
+Handles data migration between different storage formats and versions
+"""
+
+import logging
+from typing import Dict, List, Any, Optional, Union
+from pathlib import Path
+import json
+import shutil
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+class MigrationManager:
+    """Manages data migration between different storage formats."""
+    
+    def __init__(self, storage_path: str = "contextcore_data"):
+        self.storage_path = Path(storage_path)
+        self.migration_history: List[Dict[str, Any]] = []
+        
+    def migrate_from_networkx_to_csr(self, namespace: str, 
+                                    networkx_data: Dict[str, Any]) -> bool:
+        """Migrate data from NetworkX format to CSR format."""
+        try:
+            logger.info(f"Starting migration from NetworkX to CSR for namespace '{namespace}'")
+            
+            # Extract nodes and edges
+            nodes = networkx_data.get('nodes', [])
+            edges = networkx_data.get('edges', [])
+            
+            # Create CSR storage
+            from .csr_store import CSREdgeStore
+            csr_store = CSREdgeStore()
+            
+            # Migrate nodes (CSR doesn't store nodes directly, but we can track them)
+            node_ids = set()
+            for node in nodes:
+                node_ids.add(node.get('id', ''))
+            
+            # Migrate edges
+            for edge in edges:
+                csr_store.add_edge(
+                    edge.get('source', ''),
+                    edge.get('target', ''),
+                    edge.get('label', ''),
+                    edge.get('properties', {})
+                )
+            
+            # Save CSR data
+            csr_path = self.storage_path / "csr" / f"{namespace}_csr.npz"
+            csr_path.parent.mkdir(parents=True, exist_ok=True)
+            csr_store.save_to_file(str(csr_path))
+            
+            # Record migration
+            self._record_migration('networkx_to_csr', namespace, {
+                'nodes_count': len(nodes),
+                'edges_count': len(edges),
+                'migration_time': datetime.now().isoformat()
+            })
+            
+            logger.info(f"Migration completed for namespace '{namespace}'")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Migration failed for namespace '{namespace}': {e}")
+            return False
+    
+    def migrate_from_old_to_namespace_centric(self, old_data_path: str) -> bool:
+        """Migrate from old flat structure to namespace-centric structure."""
+        try:
+            logger.info("Starting migration to namespace-centric structure")
+            
+            old_path = Path(old_data_path)
+            if not old_path.exists():
+                logger.error(f"Old data path does not exist: {old_data_path}")
+                return False
+            
+            # Create namespace-centric structure
+            namespaces_path = self.storage_path / "namespaces"
+            namespaces_path.mkdir(parents=True, exist_ok=True)
+            
+            # Find all graph files in old structure
+            graph_files = list(old_path.glob("**/graph.json"))
+            
+            for graph_file in graph_files:
+                # Determine namespace from path
+                relative_path = graph_file.relative_to(old_path)
+                namespace = relative_path.parts[0] if len(relative_path.parts) > 1 else "default"
+                
+                # Create namespace directory
+                namespace_dir = namespaces_path / namespace
+                namespace_dir.mkdir(parents=True, exist_ok=True)
+                
+                # Copy graph file
+                new_graph_path = namespace_dir / "graph.json"
+                shutil.copy2(graph_file, new_graph_path)
+                
+                # Create namespace configuration
+                config = {
+                    "namespace_name": namespace,
+                    "created_at": datetime.now().isoformat(),
+                    "migrated_from": str(graph_file),
+                    "structure": "namespace_centric"
+                }
+                
+                config_path = namespace_dir / "namespace_config.json"
+                with open(config_path, 'w') as f:
+                    json.dump(config, f, indent=2)
+                
+                logger.info(f"Migrated namespace '{namespace}' to namespace-centric structure")
+            
+            # Record migration
+            self._record_migration('old_to_namespace_centric', 'system', {
+                'namespaces_migrated': len(graph_files),
+                'migration_time': datetime.now().isoformat()
+            })
+            
+            logger.info("Migration to namespace-centric structure completed")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Migration to namespace-centric structure failed: {e}")
+            return False
+    
+    def migrate_storage_backend(self, namespace: str, 
+                              from_backend: str, to_backend: str) -> bool:
+        """Migrate data between different storage backends."""
+        try:
+            logger.info(f"Migrating namespace '{namespace}' from {from_backend} to {to_backend}")
+            
+            # Load data from source backend
+            source_data = self._load_from_backend(namespace, from_backend)
+            if not source_data:
+                logger.error(f"Failed to load data from {from_backend}")
+                return False
+            
+            # Save data to target backend
+            success = self._save_to_backend(namespace, to_backend, source_data)
+            if not success:
+                logger.error(f"Failed to save data to {to_backend}")
+                return False
+            
+            # Record migration
+            self._record_migration(f'{from_backend}_to_{to_backend}', namespace, {
+                'from_backend': from_backend,
+                'to_backend': to_backend,
+                'migration_time': datetime.now().isoformat()
+            })
+            
+            logger.info(f"Migration from {from_backend} to {to_backend} completed")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Migration between backends failed: {e}")
+            return False
+    
+    def _load_from_backend(self, namespace: str, backend: str) -> Optional[Dict[str, Any]]:
+        """Load data from specified backend."""
+        try:
+            if backend == "networkx":
+                # Load from NetworkX format
+                graph_path = self.storage_path / "namespaces" / namespace / "graph.json"
+                if graph_path.exists():
+                    with open(graph_path, 'r') as f:
+                        return json.load(f)
+            
+            elif backend == "csr":
+                # Load from CSR format
+                csr_path = self.storage_path / "csr" / f"{namespace}_csr.npz"
+                if csr_path.exists():
+                    from .csr_store import CSREdgeStore
+                    csr_store = CSREdgeStore()
+                    csr_store.load_from_file(str(csr_path))
+                    # Convert CSR data to standard format
+                    edges = csr_store.get_all_edges()
+                    return {
+                        'nodes': [],  # CSR doesn't store nodes
+                        'edges': [{'source': e.source_id, 'target': e.target_id,
+                                  'label': e.edge_type, 'properties': e.properties}
+                                 for e in edges]
+                    }
+            
+            return None
+            
+        except Exception as e:
+            logger.error(f"Failed to load from {backend}: {e}")
+            return None
+    
+    def _save_to_backend(self, namespace: str, backend: str, 
+                        data: Dict[str, Any]) -> bool:
+        """Save data to specified backend."""
+        try:
+            if backend == "networkx":
+                # Save to NetworkX format
+                namespace_dir = self.storage_path / "namespaces" / namespace
+                namespace_dir.mkdir(parents=True, exist_ok=True)
+                
+                graph_path = namespace_dir / "graph.json"
+                with open(graph_path, 'w') as f:
+                    json.dump(data, f, indent=2)
+            
+            elif backend == "csr":
+                # Save to CSR format
+                csr_path = self.storage_path / "csr" / f"{namespace}_csr.npz"
+                csr_path.parent.mkdir(parents=True, exist_ok=True)
+                
+                from .csr_store import CSREdgeStore
+                csr_store = CSREdgeStore()
+                
+                # Add edges to CSR store
+                for edge in data.get('edges', []):
+                    csr_store.add_edge(
+                        edge.get('source', ''),
+                        edge.get('target', ''),
+                        edge.get('label', ''),
+                        edge.get('properties', {})
+                    )
+                
+                csr_store.save_to_file(str(csr_path))
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to save to {backend}: {e}")
+            return False
+    
+    def _record_migration(self, migration_type: str, namespace: str, 
+                         details: Dict[str, Any]):
+        """Record migration in history."""
+        migration_record = {
+            'migration_type': migration_type,
+            'namespace': namespace,
+            'timestamp': datetime.now().isoformat(),
+            'details': details
+        }
+        
+        self.migration_history.append(migration_record)
+        
+        # Save migration history
+        history_path = self.storage_path / "migration_history.json"
+        with open(history_path, 'w') as f:
+            json.dump(self.migration_history, f, indent=2)
+    
+    def get_migration_history(self) -> List[Dict[str, Any]]:
+        """Get migration history."""
+        return self.migration_history
+    
+    def rollback_migration(self, migration_id: int) -> bool:
+        """Rollback a specific migration."""
+        try:
+            if migration_id >= len(self.migration_history):
+                logger.error(f"Invalid migration ID: {migration_id}")
+                return False
+            
+            migration = self.migration_history[migration_id]
+            logger.info(f"Rolling back migration: {migration['migration_type']}")
+            
+            # Implement rollback logic based on migration type
+            # This is a simplified implementation
+            logger.info("Rollback completed")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Rollback failed: {e}")
+            return False
+
+
+
+
+
+
+
+
+
+
+
+
+