contextsynapse 1.0.0__py3-none-any.whl

contextsynapse/storage/storage_manager.py

@@ -0,0 +1,214 @@
+"""Unified storage manager supporting multiple modes."""
+
+import logging
+from pathlib import Path
+from typing import Dict, Any, Optional
+from ..core.graph_structures import GraphNode, GraphEdge
+from ..config.mode_config import NamespaceConfig, StorageMode
+from .pure_graph_storage import PureGraphStorage
+from ..metadata.metadata_db import MetadataDB
+
+logger = logging.getLogger(__name__)
+
+
+class StorageManager:
+    """Unified storage interface supporting multiple modes.
+    
+    This abstraction allows the same AIQL queries to work with different
+    storage backends (pure graph, collection-based, hybrid) without
+    changing the query syntax.
+    """
+    
+    def __init__(self, namespace: str, config: Optional[Dict[str, Any]] = None):
+        """Initialize storage manager.
+        
+        Args:
+            namespace: Namespace name
+            config: Optional configuration dictionary (if None, loads from config file)
+        """
+        self.namespace = namespace
+        self.namespace_config = NamespaceConfig(namespace)
+        
+        # Override config if provided
+        if config:
+            self.namespace_config.config.update(config)
+        
+        self.mode = self.namespace_config.mode
+        self.base_path = Path(f"contextcore_data/namespaces/{namespace}")
+        
+        # Initialize metadata database
+        self.metadata_db = None
+        try:
+            metadata_config = self.namespace_config.metadata_config
+            self.metadata_db = MetadataDB(metadata_config)
+        except Exception as e:
+            logger.warning(f"[EMOJI][EMOJI] Failed to initialize metadata DB: {e}")
+        
+        # Initialize storage backends based on mode
+        self.pure_graph_storage = None
+        self.collection_storage = None  # Will use existing _store_node_in_collection
+        
+        if self.mode == StorageMode.PURE_GRAPH:
+            self.pure_graph_storage = PureGraphStorage(namespace, self.base_path)
+        elif self.mode == StorageMode.HYBRID:
+            self.pure_graph_storage = PureGraphStorage(namespace, self.base_path)
+            # Collection storage handled by existing methods
+        
+        logger.info(f"[EMOJI] StorageManager initialized: mode={self.mode.value}, namespace={namespace}")
+    
+    def store_node(self, node: GraphNode, collection: Optional[str] = None) -> str:
+        """Store node based on configured mode.
+        
+        This method abstracts the storage implementation so AIQL queries
+        don't need to know about the underlying storage mode.
+        
+        Args:
+            node: GraphNode to store
+            collection: Collection name (used for collection-based mode)
+        
+        Returns:
+            Path where node was stored
+        """
+        node_type = node.label if hasattr(node, 'label') else 'Node'
+        stored_path = None
+        
+        try:
+            if self.mode == StorageMode.PURE_GRAPH:
+                # Store in pure graph structure
+                stored_path = self.pure_graph_storage.store_node(node)
+                
+            elif self.mode == StorageMode.COLLECTION_BASED:
+                # Store in collection-based structure (use existing method)
+                # This will be handled by the executor's existing _store_node_in_collection
+                # We return a placeholder path - actual storage happens in executor
+                stored_path = f"collections/{collection}/nodes/{node_type}/{node.id}.json"
+                
+            elif self.mode == StorageMode.HYBRID:
+                # Store in both places
+                self.pure_graph_storage.store_node(node)
+                stored_path = f"collections/{collection}/nodes/{node_type}/{node.id}.json"
+            
+            # Store metadata in SQL database
+            if self.metadata_db:
+                try:
+                    node_properties = node.properties if hasattr(node, 'properties') else {}
+                    self.metadata_db.store_metadata(
+                        entity_id=node.id,
+                        namespace=self.namespace,
+                        collection=collection,
+                        node_type=node_type,
+                        entity_type='node',
+                        metadata=node_properties
+                    )
+                except Exception as e:
+                    logger.warning(f"[EMOJI][EMOJI] Failed to store metadata for node {node.id}: {e}")
+            
+            return stored_path or f"{node_type}/{node.id}.json"
+            
+        except Exception as e:
+            logger.error(f"[EMOJI] Failed to store node {node.id}: {e}")
+            raise
+    
+    def store_edge(self, edge: GraphEdge, collection: Optional[str] = None) -> str:
+        """Store edge based on configured mode.
+        
+        Args:
+            edge: GraphEdge to store
+            collection: Collection name (used for collection-based mode)
+        
+        Returns:
+            Path where edge was stored
+        """
+        stored_path = None
+        
+        try:
+            if self.mode == StorageMode.PURE_GRAPH:
+                # Store in pure graph structure
+                stored_path = self.pure_graph_storage.store_edge(edge)
+                
+            elif self.mode == StorageMode.COLLECTION_BASED:
+                # Store in collection-based structure
+                # For graph collection, store in centralized graph/edges/
+                if collection == 'graph' or not collection:
+                    stored_path = f"collections/graph/edges/{edge.label}/{edge.id}.json"
+                else:
+                    stored_path = f"collections/{collection}/edges/{edge.label}/{edge.id}.json"
+                
+            elif self.mode == StorageMode.HYBRID:
+                # Store in pure graph AND collection-based
+                self.pure_graph_storage.store_edge(edge)
+                stored_path = f"collections/graph/edges/{edge.label}/{edge.id}.json"
+            
+            # Store metadata in SQL database
+            if self.metadata_db:
+                try:
+                    edge_properties = edge.properties if hasattr(edge, 'properties') else {}
+                    self.metadata_db.store_metadata(
+                        entity_id=edge.id,
+                        namespace=self.namespace,
+                        collection=collection or 'graph',
+                        node_type=edge.label,
+                        entity_type='edge',
+                        metadata=edge_properties
+                    )
+                except Exception as e:
+                    logger.warning(f"[EMOJI][EMOJI] Failed to store metadata for edge {edge.id}: {e}")
+            
+            return stored_path or f"edges/{edge.label}/{edge.id}.json"
+            
+        except Exception as e:
+            logger.error(f"[EMOJI] Failed to store edge {edge.id}: {e}")
+            raise
+    
+    def get_statistics(self, collection: Optional[str] = None) -> Dict[str, Any]:
+        """Get statistics from SQL metadata database.
+        
+        Args:
+            collection: Optional collection name
+        
+        Returns:
+            Statistics dictionary
+        """
+        if self.metadata_db:
+            return self.metadata_db.get_statistics(self.namespace, collection)
+        return {}
+    
+    def update_statistics(self, collection: str, node_count: int = 0, edge_count: int = 0):
+        """Update statistics in SQL metadata database.
+        
+        Args:
+            collection: Collection name
+            node_count: Number of nodes to add
+            edge_count: Number of edges to add
+        """
+        if self.metadata_db:
+            self.metadata_db.update_statistics(
+                self.namespace, collection, node_count, edge_count
+            )
+    
+    def record_query(self, query_id: str, query_text: str, query_type: Optional[str],
+                    execution_time_ms: float, rows_returned: int,
+                    metadata: Optional[Dict] = None):
+        """Record query performance metrics.
+        
+        Args:
+            query_id: Unique query identifier
+            query_text: Query text
+            query_type: Type of query
+            execution_time_ms: Execution time in milliseconds
+            rows_returned: Number of rows returned
+            metadata: Additional metadata
+        """
+        if self.metadata_db:
+            self.metadata_db.record_query(
+                query_id, query_text, self.namespace, query_type,
+                execution_time_ms, rows_returned, metadata
+            )
+    
+    def close(self):
+        """Close storage manager and connections."""
+        if self.metadata_db:
+            self.metadata_db.close()
+
+
+

contextsynapse/storage/storage_query_optimizer.py

@@ -0,0 +1,314 @@
+"""
+Query Optimizer for AIContextDB
+Optimizes queries for better performance across different storage backends
+"""
+
+import logging
+from typing import Dict, List, Any, Optional, Union, Tuple
+from dataclasses import dataclass
+from enum import Enum
+import time
+
+logger = logging.getLogger(__name__)
+
+class QueryType(Enum):
+    """Types of queries."""
+    SELECT = "select"
+    COUNT = "count"
+    TRAVERSE = "traverse"
+    AGGREGATE = "aggregate"
+    SEARCH = "search"
+
+class StorageBackend(Enum):
+    """Available storage backends."""
+    NETWORKX = "networkx"
+    CSR = "csr"
+    COLUMNAR = "columnar"
+    HYBRID = "hybrid"
+
+@dataclass
+class QueryPlan:
+    """Represents an optimized query plan."""
+    query_type: QueryType
+    storage_backend: StorageBackend
+    estimated_cost: float
+    estimated_time: float
+    operations: List[str]
+    indexes_used: List[str]
+    optimization_applied: List[str]
+
+class QueryOptimizer:
+    """Optimizes queries for different storage backends."""
+    
+    def __init__(self):
+        self.optimization_rules = {
+            QueryType.SELECT: self._optimize_select_query,
+            QueryType.COUNT: self._optimize_count_query,
+            QueryType.TRAVERSE: self._optimize_traverse_query,
+            QueryType.AGGREGATE: self._optimize_aggregate_query,
+            QueryType.SEARCH: self._optimize_search_query
+        }
+        
+        # Performance statistics
+        self.stats = {
+            'total_queries': 0,
+            'optimized_queries': 0,
+            'cache_hits': 0,
+            'average_optimization_time': 0.0
+        }
+    
+    def optimize_query(self, query: str, query_type: QueryType,
+                      data_stats: Dict[str, Any]) -> QueryPlan:
+        """Optimize a query based on its type and data characteristics."""
+        start_time = time.time()
+        
+        try:
+            # Determine best storage backend
+            backend = self._select_optimal_backend(query_type, data_stats)
+            
+            # Apply query-specific optimizations
+            optimization_func = self.optimization_rules.get(query_type)
+            if optimization_func:
+                operations, indexes, optimizations = optimization_func(query, data_stats)
+            else:
+                operations, indexes, optimizations = [query], [], []
+            
+            # Estimate cost and time
+            cost = self._estimate_cost(query_type, data_stats, backend)
+            estimated_time = self._estimate_time(query_type, data_stats, backend)
+            
+            # Create query plan
+            plan = QueryPlan(
+                query_type=query_type,
+                storage_backend=backend,
+                estimated_cost=cost,
+                estimated_time=estimated_time,
+                operations=operations,
+                indexes_used=indexes,
+                optimization_applied=optimizations
+            )
+            
+            # Update statistics
+            self._update_stats(time.time() - start_time)
+            
+            logger.info(f"Query optimized: {query_type.value} -> {backend.value}")
+            return plan
+            
+        except Exception as e:
+            logger.error(f"Query optimization failed: {e}")
+            # Return default plan
+            return QueryPlan(
+                query_type=query_type,
+                storage_backend=StorageBackend.NETWORKX,
+                estimated_cost=1.0,
+                estimated_time=0.1,
+                operations=[query],
+                indexes_used=[],
+                optimization_applied=[]
+            )
+    
+    def _select_optimal_backend(self, query_type: QueryType, 
+                              data_stats: Dict[str, Any]) -> StorageBackend:
+        """Select the optimal storage backend for the query."""
+        node_count = data_stats.get('node_count', 0)
+        edge_count = data_stats.get('edge_count', 0)
+        
+        if query_type == QueryType.TRAVERSE:
+            # CSR is optimal for traversal queries
+            if edge_count > 1000:
+                return StorageBackend.CSR
+            else:
+                return StorageBackend.NETWORKX
+        
+        elif query_type == QueryType.COUNT:
+            # Columnar is optimal for count queries
+            if node_count > 10000:
+                return StorageBackend.COLUMNAR
+            else:
+                return StorageBackend.NETWORKX
+        
+        elif query_type == QueryType.AGGREGATE:
+            # Columnar is optimal for aggregations
+            return StorageBackend.COLUMNAR
+        
+        elif query_type == QueryType.SEARCH:
+            # Hybrid is optimal for search queries
+            return StorageBackend.HYBRID
+        
+        else:
+            # Default to NetworkX for simple queries
+            return StorageBackend.NETWORKX
+    
+    def _optimize_select_query(self, query: str, data_stats: Dict[str, Any]) -> Tuple[List[str], List[str], List[str]]:
+        """Optimize SELECT queries."""
+        operations = [query]
+        indexes = []
+        optimizations = []
+        
+        # Add index usage if available
+        if 'text_index' in data_stats.get('available_indexes', []):
+            indexes.append('text_index')
+            optimizations.append('text_index_usage')
+        
+        # Add projection optimization
+        if 'WHERE' in query.upper():
+            optimizations.append('predicate_pushdown')
+        
+        return operations, indexes, optimizations
+    
+    def _optimize_count_query(self, query: str, data_stats: Dict[str, Any]) -> Tuple[List[str], List[str], List[str]]:
+        """Optimize COUNT queries."""
+        operations = [query]
+        indexes = []
+        optimizations = []
+        
+        # Use columnar storage for count queries
+        optimizations.append('columnar_count')
+        
+        # Add index usage if available
+        if 'count_index' in data_stats.get('available_indexes', []):
+            indexes.append('count_index')
+            optimizations.append('count_index_usage')
+        
+        return operations, indexes, optimizations
+    
+    def _optimize_traverse_query(self, query: str, data_stats: Dict[str, Any]) -> Tuple[List[str], List[str], List[str]]:
+        """Optimize TRAVERSE queries."""
+        operations = [query]
+        indexes = []
+        optimizations = []
+        
+        # Use CSR for traversal
+        optimizations.append('csr_traversal')
+        
+        # Add graph index usage
+        if 'graph_index' in data_stats.get('available_indexes', []):
+            indexes.append('graph_index')
+            optimizations.append('graph_index_usage')
+        
+        # Add depth limiting
+        if 'DEPTH' not in query.upper():
+            optimizations.append('depth_limiting')
+        
+        return operations, indexes, optimizations
+    
+    def _optimize_aggregate_query(self, query: str, data_stats: Dict[str, Any]) -> Tuple[List[str], List[str], List[str]]:
+        """Optimize AGGREGATE queries."""
+        operations = [query]
+        indexes = []
+        optimizations = []
+        
+        # Use columnar storage for aggregations
+        optimizations.append('columnar_aggregation')
+        
+        # Add grouping optimization
+        if 'GROUP BY' in query.upper():
+            optimizations.append('grouping_optimization')
+        
+        return operations, indexes, optimizations
+    
+    def _optimize_search_query(self, query: str, data_stats: Dict[str, Any]) -> Tuple[List[str], List[str], List[str]]:
+        """Optimize SEARCH queries."""
+        operations = [query]
+        indexes = []
+        optimizations = []
+        
+        # Use hybrid search
+        optimizations.append('hybrid_search')
+        
+        # Add vector index usage
+        if 'vector_index' in data_stats.get('available_indexes', []):
+            indexes.append('vector_index')
+            optimizations.append('vector_index_usage')
+        
+        # Add text index usage
+        if 'text_index' in data_stats.get('available_indexes', []):
+            indexes.append('text_index')
+            optimizations.append('text_index_usage')
+        
+        return operations, indexes, optimizations
+    
+    def _estimate_cost(self, query_type: QueryType, data_stats: Dict[str, Any], 
+                      backend: StorageBackend) -> float:
+        """Estimate the cost of executing the query."""
+        node_count = data_stats.get('node_count', 0)
+        edge_count = data_stats.get('edge_count', 0)
+        
+        # Base cost by query type
+        base_costs = {
+            QueryType.SELECT: 1.0,
+            QueryType.COUNT: 0.5,
+            QueryType.TRAVERSE: 2.0,
+            QueryType.AGGREGATE: 1.5,
+            QueryType.SEARCH: 3.0
+        }
+        
+        base_cost = base_costs.get(query_type, 1.0)
+        
+        # Scale by data size
+        if backend == StorageBackend.CSR:
+            scale_factor = edge_count / 1000.0
+        elif backend == StorageBackend.COLUMNAR:
+            scale_factor = node_count / 10000.0
+        else:
+            scale_factor = (node_count + edge_count) / 1000.0
+        
+        return base_cost * max(1.0, scale_factor)
+    
+    def _estimate_time(self, query_type: QueryType, data_stats: Dict[str, Any], 
+                      backend: StorageBackend) -> float:
+        """Estimate the execution time of the query."""
+        cost = self._estimate_cost(query_type, data_stats, backend)
+        
+        # Convert cost to estimated time (milliseconds)
+        time_multipliers = {
+            StorageBackend.NETWORKX: 10.0,
+            StorageBackend.CSR: 1.0,
+            StorageBackend.COLUMNAR: 2.0,
+            StorageBackend.HYBRID: 5.0
+        }
+        
+        multiplier = time_multipliers.get(backend, 10.0)
+        return cost * multiplier
+    
+    def _update_stats(self, optimization_time: float):
+        """Update optimization statistics."""
+        self.stats['total_queries'] += 1
+        self.stats['optimized_queries'] += 1
+        
+        # Update average optimization time
+        total_time = self.stats['average_optimization_time'] * (self.stats['optimized_queries'] - 1)
+        self.stats['average_optimization_time'] = (total_time + optimization_time) / self.stats['optimized_queries']
+    
+    def get_optimization_stats(self) -> Dict[str, Any]:
+        """Get optimization statistics."""
+        return self.stats.copy()
+    
+    def suggest_indexes(self, query_patterns: List[str]) -> List[str]:
+        """Suggest indexes based on query patterns."""
+        suggestions = []
+        
+        for pattern in query_patterns:
+            if 'WHERE' in pattern.upper():
+                suggestions.append('text_index')
+            if 'TRAVERSE' in pattern.upper():
+                suggestions.append('graph_index')
+            if 'SEARCH' in pattern.upper():
+                suggestions.append('vector_index')
+            if 'COUNT' in pattern.upper():
+                suggestions.append('count_index')
+        
+        return list(set(suggestions))
+
+
+
+
+
+
+
+
+
+
+
+
+