"""SQLite graph store for Headroom's knowledge graph memory system. Provides persistent storage for entities and relationships with efficient lookup via database indexes and BFS-based traversal. Memory usage is bounded by SQLite's page cache. This is a drop-in replacement for InMemoryGraphStore that: - Persists all data to disk - Keeps memory bounded (configurable page cache) - Maintains the same async interface - Supports all query patterns (by ID, by name, BFS traversal) """ from __future__ import annotations import json import sqlite3 from collections import deque from datetime import datetime from pathlib import Path from threading import RLock from typing import TYPE_CHECKING, Any from .graph_models import Entity, Relationship, RelationshipDirection, Subgraph if TYPE_CHECKING: from ..tracker import ComponentStats class SQLiteGraphStore: """SQLite-based graph store implementing the GraphStore protocol. Provides persistent storage for entities and relationships with efficient lookup via database indexes. All operations are thread-safe. Features: - O(log n) entity and relationship lookup by ID (indexed) - O(log n) entity lookup by name (per user, case-insensitive) - O(log n) lookup of relationships by source or target entity - BFS-based subgraph traversal with configurable hop limit - BFS-based shortest path finding between entities - Configurable page cache for memory bounding Schema: - entities: id, user_id, name, name_lower, entity_type, description, properties, created_at, updated_at, metadata - relationships: id, user_id, source_id, target_id, relation_type, weight, properties, created_at, metadata Usage: store = SQLiteGraphStore("./graph.db") await store.add_entity(Entity(user_id="alice", name="Project X", entity_type="project")) entity = await store.get_entity_by_name("alice", "project x") # Case-insensitive subgraph = await store.query_subgraph(["entity-id"], max_hops=2) """ def __init__( self, db_path: str | Path = "headroom_graph.db", page_cache_size_kb: int = 8192, # 8MB default cache ) -> None: """Initialize the SQLite graph store. Args: db_path: Path to SQLite database file. Created if it doesn't exist. page_cache_size_kb: SQLite page cache size in KB. Higher = more memory, faster queries. Set to -1 for default SQLite behavior. """ self.db_path = Path(db_path) self._page_cache_size_kb = page_cache_size_kb self._lock = RLock() self._init_db() def _get_conn(self) -> sqlite3.Connection: """Get a new database connection (thread-safe pattern). Returns: A new SQLite connection with row factory configured. """ conn = sqlite3.connect(str(self.db_path)) conn.row_factory = sqlite3.Row # Configure page cache size (negative = KB, positive = pages) if self._page_cache_size_kb > 0: conn.execute(f"PRAGMA cache_size = -{self._page_cache_size_kb}") # Enable foreign keys conn.execute("PRAGMA foreign_keys = ON") return conn def _init_db(self) -> None: """Initialize the database schema with indexes.""" with self._get_conn() as conn: # Create entities table conn.execute(""" CREATE TABLE IF NOT EXISTS entities ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL, name TEXT NOT NULL, name_lower TEXT NOT NULL, entity_type TEXT NOT NULL DEFAULT 'unknown', description TEXT, properties TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, updated_at TEXT NOT NULL, metadata TEXT NOT NULL DEFAULT '{}' ) """) # Create relationships table conn.execute(""" CREATE TABLE IF NOT EXISTS relationships ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL, source_id TEXT NOT NULL, target_id TEXT NOT NULL, relation_type TEXT NOT NULL DEFAULT 'related_to', weight REAL NOT NULL DEFAULT 1.0, properties TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, metadata TEXT NOT NULL DEFAULT '{}', FOREIGN KEY (source_id) REFERENCES entities(id) ON DELETE CASCADE, FOREIGN KEY (target_id) REFERENCES entities(id) ON DELETE CASCADE ) """) # Create indexes for efficient querying # Entity indexes conn.execute("CREATE INDEX IF NOT EXISTS idx_entities_user_id ON entities(user_id)") conn.execute( "CREATE INDEX IF NOT EXISTS idx_entities_name_lookup " "ON entities(user_id, name_lower)" ) conn.execute("CREATE INDEX IF NOT EXISTS idx_entities_type ON entities(entity_type)") # Relationship indexes conn.execute( "CREATE INDEX IF NOT EXISTS idx_relationships_source ON relationships(source_id)" ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_relationships_target ON relationships(target_id)" ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_relationships_type ON relationships(relation_type)" ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_relationships_user ON relationships(user_id)" ) conn.commit() def _entity_to_row(self, entity: Entity) -> dict[str, Any]: """Convert Entity object to row dict for insertion.""" return { "id": entity.id, "user_id": entity.user_id, "name": entity.name, "name_lower": entity.name.lower(), "entity_type": entity.entity_type, "description": entity.description, "properties": json.dumps(entity.properties), "created_at": entity.created_at.isoformat(), "updated_at": entity.updated_at.isoformat(), "metadata": json.dumps(entity.metadata), } def _row_to_entity(self, row: sqlite3.Row) -> Entity: """Convert database row to Entity object.""" return Entity( id=row["id"], user_id=row["user_id"], name=row["name"], entity_type=row["entity_type"], description=row["description"], properties=json.loads(row["properties"]), created_at=datetime.fromisoformat(row["created_at"]), updated_at=datetime.fromisoformat(row["updated_at"]), metadata=json.loads(row["metadata"]), ) def _relationship_to_row(self, relationship: Relationship) -> dict[str, Any]: """Convert Relationship object to row dict for insertion.""" return { "id": relationship.id, "user_id": relationship.user_id, "source_id": relationship.source_id, "target_id": relationship.target_id, "relation_type": relationship.relation_type, "weight": relationship.weight, "properties": json.dumps(relationship.properties), "created_at": relationship.created_at.isoformat(), "metadata": json.dumps(relationship.metadata), } def _row_to_relationship(self, row: sqlite3.Row) -> Relationship: """Convert database row to Relationship object.""" return Relationship( id=row["id"], user_id=row["user_id"], source_id=row["source_id"], target_id=row["target_id"], relation_type=row["relation_type"], weight=row["weight"], properties=json.loads(row["properties"]), created_at=datetime.fromisoformat(row["created_at"]), metadata=json.loads(row["metadata"]), ) # ========================================================================= # Entity Operations # ========================================================================= async def add_entity(self, entity: Entity) -> None: """Add an entity to the graph store. If an entity with the same ID already exists, it will be replaced. Args: entity: The entity to add. """ row = self._entity_to_row(entity) with self._lock: with self._get_conn() as conn: conn.execute( """ INSERT OR REPLACE INTO entities ( id, user_id, name, name_lower, entity_type, description, properties, created_at, updated_at, metadata ) VALUES ( :id, :user_id, :name, :name_lower, :entity_type, :description, :properties, :created_at, :updated_at, :metadata ) """, row, ) conn.commit() async def get_entity(self, entity_id: str) -> Entity | None: """Retrieve an entity by ID. Args: entity_id: The unique identifier of the entity. Returns: The entity if found, None otherwise. """ with self._lock: with self._get_conn() as conn: cursor = conn.execute( "SELECT * FROM entities WHERE id = ?", (entity_id,), ) row = cursor.fetchone() if row is None: return None return self._row_to_entity(row) async def get_entity_by_name(self, user_id: str, name: str) -> Entity | None: """Retrieve an entity by name (case-insensitive). Args: user_id: The user who owns the entity. name: The name of the entity (case-insensitive). Returns: The entity if found, None otherwise. """ with self._lock: with self._get_conn() as conn: cursor = conn.execute( "SELECT * FROM entities WHERE user_id = ? AND name_lower = ?", (user_id, name.lower()), ) row = cursor.fetchone() if row is None: return None return self._row_to_entity(row) async def delete_entity(self, entity_id: str) -> bool: """Delete an entity and all its relationships. Relationships are automatically deleted via ON DELETE CASCADE. Args: entity_id: The unique identifier of the entity. Returns: True if the entity was deleted, False if not found. """ with self._lock: with self._get_conn() as conn: cursor = conn.execute( "DELETE FROM entities WHERE id = ?", (entity_id,), ) conn.commit() return cursor.rowcount > 0 # ========================================================================= # Relationship Operations # ========================================================================= async def add_relationship(self, relationship: Relationship) -> None: """Add a relationship to the graph store. If a relationship with the same ID already exists, it will be replaced. Args: relationship: The relationship to add. """ row = self._relationship_to_row(relationship) with self._lock: with self._get_conn() as conn: conn.execute( """ INSERT OR REPLACE INTO relationships ( id, user_id, source_id, target_id, relation_type, weight, properties, created_at, metadata ) VALUES ( :id, :user_id, :source_id, :target_id, :relation_type, :weight, :properties, :created_at, :metadata ) """, row, ) conn.commit() async def get_relationships( self, entity_id: str, direction: RelationshipDirection = RelationshipDirection.BOTH, relation_type: str | None = None, ) -> list[Relationship]: """Get relationships for an entity. Args: entity_id: The entity ID to get relationships for. direction: Whether to get outgoing, incoming, or both relationships. relation_type: Optional filter for relationship type. Returns: List of matching relationships. """ with self._lock: with self._get_conn() as conn: conditions = [] params: list[Any] = [] if direction == RelationshipDirection.OUTGOING: conditions.append("source_id = ?") params.append(entity_id) elif direction == RelationshipDirection.INCOMING: conditions.append("target_id = ?") params.append(entity_id) else: # BOTH conditions.append("(source_id = ? OR target_id = ?)") params.extend([entity_id, entity_id]) if relation_type is not None: conditions.append("relation_type = ?") params.append(relation_type) where_clause = " AND ".join(conditions) cursor = conn.execute( f"SELECT * FROM relationships WHERE {where_clause}", # nosec B608 params, ) return [self._row_to_relationship(row) for row in cursor] async def delete_relationship(self, relationship_id: str) -> bool: """Delete a single relationship. Args: relationship_id: The unique identifier of the relationship. Returns: True if the relationship was deleted, False if not found. """ with self._lock: with self._get_conn() as conn: cursor = conn.execute( "DELETE FROM relationships WHERE id = ?", (relationship_id,), ) conn.commit() return cursor.rowcount > 0 # ========================================================================= # Graph Traversal Operations # ========================================================================= async def query_subgraph( self, entity_ids: list[str], max_hops: int = 2, direction: RelationshipDirection = RelationshipDirection.BOTH, relation_types: list[str] | None = None, ) -> Subgraph: """Query a subgraph starting from given entities using BFS traversal. Performs a breadth-first traversal from the starting entities, collecting all entities and relationships within the specified number of hops. Args: entity_ids: Starting entity IDs for the traversal. max_hops: Maximum number of hops from starting entities (default 2). direction: Direction of relationship traversal. relation_types: Optional filter for relationship types. Returns: Subgraph containing all reachable entities and relationships. """ with self._lock: with self._get_conn() as conn: collected_entities: dict[str, Entity] = {} collected_relationships: dict[str, Relationship] = {} # BFS queue: (entity_id, current_depth) queue: deque[tuple[str, int]] = deque() visited: set[str] = set() # Initialize queue with starting entities for entity_id in entity_ids: cursor = conn.execute( "SELECT * FROM entities WHERE id = ?", (entity_id,), ) row = cursor.fetchone() if row is not None: queue.append((entity_id, 0)) visited.add(entity_id) collected_entities[entity_id] = self._row_to_entity(row) # BFS traversal while queue: current_id, depth = queue.popleft() if depth >= max_hops: continue # Build relationship query based on direction if direction == RelationshipDirection.OUTGOING: rel_query = "SELECT * FROM relationships WHERE source_id = ?" rel_params: list[Any] = [current_id] elif direction == RelationshipDirection.INCOMING: rel_query = "SELECT * FROM relationships WHERE target_id = ?" rel_params = [current_id] else: # BOTH rel_query = ( "SELECT * FROM relationships WHERE source_id = ? OR target_id = ?" ) rel_params = [current_id, current_id] # Filter by relation types if specified if relation_types is not None and len(relation_types) > 0: placeholders = ", ".join("?" * len(relation_types)) rel_query += f" AND relation_type IN ({placeholders})" rel_params.extend(relation_types) cursor = conn.execute(rel_query, rel_params) for rel_row in cursor: rel = self._row_to_relationship(rel_row) # Add relationship collected_relationships[rel.id] = rel # Determine neighbor based on direction neighbor_id = None if direction == RelationshipDirection.OUTGOING: if rel.source_id == current_id: neighbor_id = rel.target_id elif direction == RelationshipDirection.INCOMING: if rel.target_id == current_id: neighbor_id = rel.source_id else: # BOTH if rel.source_id == current_id: neighbor_id = rel.target_id elif rel.target_id == current_id: neighbor_id = rel.source_id if neighbor_id is not None and neighbor_id not in visited: # Fetch neighbor entity neighbor_cursor = conn.execute( "SELECT * FROM entities WHERE id = ?", (neighbor_id,), ) neighbor_row = neighbor_cursor.fetchone() if neighbor_row is not None: visited.add(neighbor_id) collected_entities[neighbor_id] = self._row_to_entity(neighbor_row) queue.append((neighbor_id, depth + 1)) return Subgraph( entities=list(collected_entities.values()), relationships=list(collected_relationships.values()), root_entity_ids=entity_ids, ) async def find_path( self, source_id: str, target_id: str, max_depth: int = 10, direction: RelationshipDirection = RelationshipDirection.BOTH, ) -> list[str] | None: """Find the shortest path between two entities using BFS. Args: source_id: Starting entity ID. target_id: Target entity ID. max_depth: Maximum path length to search (default 10). direction: Direction of relationship traversal. Returns: List of entity IDs representing the path (including source and target), or None if no path exists within max_depth. """ with self._lock: with self._get_conn() as conn: # Edge cases if source_id == target_id: cursor = conn.execute( "SELECT id FROM entities WHERE id = ?", (source_id,), ) return [source_id] if cursor.fetchone() else None # Check both exist cursor = conn.execute( "SELECT id FROM entities WHERE id IN (?, ?)", (source_id, target_id), ) found_ids = {row["id"] for row in cursor} if source_id not in found_ids or target_id not in found_ids: return None # BFS with path tracking queue: deque[tuple[str, list[str]]] = deque() visited: set[str] = set() queue.append((source_id, [source_id])) visited.add(source_id) while queue: current_id, path = queue.popleft() if len(path) > max_depth: continue # Build relationship query if direction == RelationshipDirection.OUTGOING: rel_query = "SELECT * FROM relationships WHERE source_id = ?" rel_params = [current_id] elif direction == RelationshipDirection.INCOMING: rel_query = "SELECT * FROM relationships WHERE target_id = ?" rel_params = [current_id] else: rel_query = ( "SELECT * FROM relationships WHERE source_id = ? OR target_id = ?" ) rel_params = [current_id, current_id] cursor = conn.execute(rel_query, rel_params) for rel_row in cursor: # Determine neighbor neighbor_id = None if direction == RelationshipDirection.OUTGOING: if rel_row["source_id"] == current_id: neighbor_id = rel_row["target_id"] elif direction == RelationshipDirection.INCOMING: if rel_row["target_id"] == current_id: neighbor_id = rel_row["source_id"] else: if rel_row["source_id"] == current_id: neighbor_id = rel_row["target_id"] elif rel_row["target_id"] == current_id: neighbor_id = rel_row["source_id"] if neighbor_id is None or neighbor_id in visited: continue new_path = path + [neighbor_id] if neighbor_id == target_id: return new_path if len(new_path) <= max_depth: visited.add(neighbor_id) queue.append((neighbor_id, new_path)) return None # ========================================================================= # User Management Operations # ========================================================================= async def clear_user(self, user_id: str) -> tuple[int, int]: """Clear all entities and relationships for a user. Args: user_id: The user ID to clear data for. Returns: Tuple of (entities_deleted, relationships_deleted). """ with self._lock: with self._get_conn() as conn: # Delete relationships for user first (also cascade-deleted) cursor = conn.execute( "DELETE FROM relationships WHERE user_id = ?", (user_id,), ) relationships_deleted = cursor.rowcount # Delete entities (cascades remaining relationships) cursor = conn.execute( "DELETE FROM entities WHERE user_id = ?", (user_id,), ) entities_deleted = cursor.rowcount conn.commit() return entities_deleted, relationships_deleted # ========================================================================= # Utility Methods # ========================================================================= async def get_entities_for_user(self, user_id: str) -> list[Entity]: """Get all entities for a user. Args: user_id: The user ID to get entities for. Returns: List of all entities belonging to the user. """ with self._lock: with self._get_conn() as conn: cursor = conn.execute( "SELECT * FROM entities WHERE user_id = ?", (user_id,), ) return [self._row_to_entity(row) for row in cursor] async def clear(self) -> None: """Clear all data from the store.""" with self._lock: with self._get_conn() as conn: conn.execute("DELETE FROM relationships") conn.execute("DELETE FROM entities") conn.commit() @property def entity_count(self) -> int: """Get the total number of entities.""" with self._lock: with self._get_conn() as conn: cursor = conn.execute("SELECT COUNT(*) FROM entities") result = cursor.fetchone()[0] return int(result) @property def relationship_count(self) -> int: """Get the total number of relationships.""" with self._lock: with self._get_conn() as conn: cursor = conn.execute("SELECT COUNT(*) FROM relationships") result = cursor.fetchone()[0] return int(result) def stats(self) -> dict: """Get store statistics. Returns: Dict with counts and database info. """ with self._lock: with self._get_conn() as conn: entity_count = conn.execute("SELECT COUNT(*) FROM entities").fetchone()[0] rel_count = conn.execute("SELECT COUNT(*) FROM relationships").fetchone()[0] users_count = conn.execute( "SELECT COUNT(DISTINCT user_id) FROM entities" ).fetchone()[0] # Get database file size db_size = self.db_path.stat().st_size if self.db_path.exists() else 0 return { "entity_count": entity_count, "relationship_count": rel_count, "users_count": users_count, "db_path": str(self.db_path), "db_size_bytes": db_size, "page_cache_size_kb": self._page_cache_size_kb, } def get_memory_stats(self) -> ComponentStats: """Get memory statistics for the MemoryTracker. Note: SQLite manages its own memory via page cache. We report the configured cache size plus estimated Python overhead. Returns: ComponentStats with current memory usage. """ import sys from ..tracker import ComponentStats with self._lock: with self._get_conn() as conn: entity_count = conn.execute("SELECT COUNT(*) FROM entities").fetchone()[0] rel_count = conn.execute("SELECT COUNT(*) FROM relationships").fetchone()[0] # Estimate Python overhead (connection objects, etc.) # The actual data is on disk, managed by SQLite's page cache python_overhead = sys.getsizeof(self) + sys.getsizeof(self._lock) # Page cache memory (this is the bounded amount) page_cache_bytes = self._page_cache_size_kb * 1024 return ComponentStats( name="sqlite_graph_store", entry_count=entity_count + rel_count, size_bytes=python_overhead + page_cache_bytes, budget_bytes=page_cache_bytes, # Cache size is the budget hits=0, misses=0, evictions=0, ) def vacuum(self) -> None: """Reclaim unused space in the database file. Call this periodically after many deletes to reduce file size. """ with self._lock: with self._get_conn() as conn: conn.execute("VACUUM") def close(self) -> None: """Close any open connections (cleanup). Note: This store uses connection-per-request pattern, so there's typically nothing to close. """ pass