MCP KQL Server

""" KQL Query Validator Pre-execution validation to catch errors before sending queries to Kusto. Validates column names, table names, and operator syntax against schema. Author: Arjun Trivedi Email: arjuntrivedi42@yahoo.com """ import re import logging from typing import Dict, List, Any logger = logging.getLogger(__name__) class KQLValidator: """Validates KQL queries against schema before execution.""" def __init__(self, memory_manager, schema_manager): """ Initialize validator with memory and schema managers. Args: memory_manager: MemoryManager instance for schema lookup schema_manager: SchemaManager instance for live schema discovery """ self.memory = memory_manager self.schema_manager = schema_manager async def validate_query( self, query: str, cluster: str, database: str, auto_discover: bool = True ) -> Dict[str, Any]: """ Validate KQL query against schema. Args: query: KQL query to validate cluster: Cluster URL database: Database name auto_discover: Whether to auto-discover missing schemas Returns: Dict with validation results: { "valid": bool, "errors": List[str], "warnings": List[str], "suggestions": List[str], "tables_used": List[str], "columns_validated": int } """ errors = [] warnings = [] suggestions = [] try: # Extract tables from query tables = self._extract_tables(query) logger.debug("Extracted tables from query: %s", tables) if not tables: warnings.append("No tables detected in query") return { "valid": True, "errors": [], "warnings": warnings, "suggestions": ["Ensure your query references at least one table"], "tables_used": [], "columns_validated": 0 } # Ensure schemas exist for all tables if auto_discover: await self._ensure_schemas_exist(cluster, database, tables) # Validate each table and its columns columns_validated = 0 for table in tables: table_errors, table_warnings, validated_count = await self._validate_table_usage( query, table, cluster, database ) errors.extend(table_errors) warnings.extend(table_warnings) columns_validated += validated_count # Validate operator syntax syntax_errors = self._validate_operator_syntax(query) errors.extend(syntax_errors) # Generate suggestions based on errors if errors: suggestions = self._generate_suggestions(errors, tables) return { "valid": len(errors) == 0, "errors": errors, "warnings": warnings, "suggestions": suggestions, "tables_used": tables, "columns_validated": columns_validated } except Exception as e: logger.error("Validation error: %s", e, exc_info=True) return { "valid": False, "errors": [f"Validation failed: {str(e)}"], "warnings": [], "suggestions": ["Check query syntax and try again"], "tables_used": [], "columns_validated": 0 } def _extract_tables(self, query: str) -> List[str]: """Extract table names from KQL query.""" tables = set() # Pattern 1: Table at start or after pipe # Matches: TableName | ..., | TableName, etc. pattern1 = r'(?:^|\|\s*)([A-Za-z_][\w]*)\s*(?:\||$)' # Pattern 2: Join/union operations # Matches: join Table, union Table, lookup Table pattern2 = r'(?:join|union|lookup)\s+(?:kind\s*=\s*\w+\s+)?([A-Za-z_][\w]*)' # Pattern 3: Bracketed table names # Matches: ['TableName'], ["TableName"] pattern3 = r'\[[\'"]([\ w\s-]+)[\'"]\]' for pattern in [pattern1, pattern2, pattern3]: matches = re.finditer(pattern, query, re.IGNORECASE) for match in matches: table_name = match.group(1).strip() # Filter out KQL keywords if not self._is_kql_keyword(table_name): tables.add(table_name) return list(tables) def _is_kql_keyword(self, word: str) -> bool: """Check if word is a KQL keyword.""" keywords = { 'where', 'project', 'extend', 'summarize', 'take', 'limit', 'sort', 'order', 'top', 'count', 'distinct', 'join', 'union', 'let', 'as', 'by', 'on', 'and', 'or', 'not', 'in', 'has', 'contains', 'startswith', 'endswith', 'between', 'ago', 'now', 'datetime', 'timespan', 'render', 'evaluate', 'parse', 'mv-expand' } return word.lower() in keywords async def _ensure_schemas_exist( self, cluster: str, database: str, tables: List[str] ) -> None: """Ensure schemas exist for all tables, trigger discovery if needed.""" for table in tables: try: schemas = self.memory._get_database_schema(cluster, database) schema = next((s for s in schemas if s.get("table") == table), None) if not schema or not schema.get("columns"): logger.info("Auto-discovering schema for %s", table) await self.schema_manager.get_table_schema( cluster=cluster, database=database, table=table, _force_refresh=False ) except Exception as e: logger.warning("Schema discovery failed for %s: %s", table, e) async def _validate_table_usage( self, query: str, table: str, cluster: str, database: str ) -> tuple[List[str], List[str], int]: """ Validate column usage for a specific table. Returns: (errors, warnings, validated_count) """ errors: List[str] = [] warnings = [] validated_count = 0 try: # Get schema for table schemas = self.memory._get_database_schema(cluster, database) schema = next((s for s in schemas if s.get("table") == table), None) if not schema: warnings.append(f"Schema not found for table '{table}'") return errors, warnings, 0 columns = schema.get("columns", {}) if not columns: errors.append(f"No columns found in schema for '{table}'. Cannot validate query.") return errors, warnings, 0 # Extract column references for this table column_refs = self._extract_column_references(query, table) # Validate each column reference for col_ref in column_refs: if col_ref not in columns: errors.append( f"Column '{col_ref}' not found in table '{table}'. " f"Available columns: {', '.join(sorted(columns.keys()))}" ) else: validated_count += 1 return errors, warnings, validated_count except Exception as e: logger.warning("Table validation failed for %s: %s", table, e) return errors, warnings, 0 def _extract_column_references(self, query: str, _table: str) -> List[str]: """Extract column references from query.""" columns = set() # Pattern for column references in common operators # Matches: where Column, project Column, extend Column, summarize by Column patterns = [ r'(?:where|project|extend|summarize|by)\s+([A-Za-z_][\w]*)', r'([A-Za-z_][\w]*)\s*(?:==|!=|<|>|<=|>=|contains|has|startswith)', ] for pattern in patterns: matches = re.finditer(pattern, query, re.IGNORECASE) for match in matches: col_name = match.group(1).strip() if not self._is_kql_keyword(col_name): columns.add(col_name) return list(columns) def _validate_operator_syntax(self, query: str) -> List[str]: """Validate KQL operator syntax.""" errors = [] # Check for invalid negation operators with spaces invalid_patterns = [ (r'!\s+=', "Invalid negation '! ='. Use '!=' without space"), (r'!\s+contains', "Invalid negation '! contains'. Use '!contains' without space"), (r'!\s+in\b', "Invalid negation '! in'. Use '!in' without space"), (r'!\s+has\b', "Invalid negation '! has'. Use '!has' without space"), (r'!has_any', "Invalid operator '!has_any'. Use '!in' instead"), ] for pattern, error_msg in invalid_patterns: if re.search(pattern, query, re.IGNORECASE): errors.append(error_msg) return errors def _generate_suggestions(self, errors: List[str], tables: List[str]) -> List[str]: """Generate helpful suggestions based on errors.""" suggestions = [] if any("not found in table" in err for err in errors): suggestions.append("Check column names against the schema") suggestions.append("Use schema discovery to refresh table schemas") if any("negation" in err.lower() for err in errors): suggestions.append("Remove spaces in negation operators (use !=, !contains, !in, !has)") if any("!has_any" in err for err in errors): suggestions.append("Replace !has_any with !in operator") if not suggestions: suggestions.append(f"Review the query syntax for tables: {', '.join(tables)}") return suggestions