ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/tools/db.py

@@ -0,0 +1,1800 @@
+"""
+Unified Database Tool for AI-Parrot
+
+Consolidates schema extraction, knowledge base building, query generation,
+validation, and execution into a single, powerful database interface.
+"""
+from __future__ import annotations
+from typing import Dict, List, Optional, Any, Union, Literal, Tuple
+import re
+import asyncio
+import json
+import hashlib
+from datetime import datetime, timedelta, timezone
+from enum import Enum
+import pandas as pd
+from pydantic import (
+    BaseModel,
+    Field,
+    field_validator,
+    model_validator
+)
+from asyncdb import AsyncDB
+from .abstract import (
+    AbstractTool,
+    ToolResult,
+    AbstractToolArgsSchema
+)
+from ..stores.abstract import AbstractStore
+from ..clients.base import AbstractClient
+from ..clients.factory import LLMFactory
+from ..models import AIMessage
+
+
+class DatabaseFlavor(str, Enum):
+    """Supported database flavors."""
+    POSTGRESQL = "postgresql"
+    MYSQL = "mysql"
+    SQLSERVER = "sqlserver"
+    BIGQUERY = "bigquery"
+    INFLUXDB = "influxdb"
+    CASSANDRA = "cassandra"
+    MONGODB = "mongodb"
+    ELASTICSEARCH = "elasticsearch"
+    SQLITE = "sqlite"
+    DUCKDB = "duckdb"
+
+
+class QueryType(str, Enum):
+    """Supported query types."""
+    SELECT = "SELECT"
+    INSERT = "INSERT"
+    UPDATE = "UPDATE"
+    DELETE = "DELETE"
+    CREATE = "CREATE"
+    ALTER = "ALTER"
+    DROP = "DROP"
+
+
+class OutputFormat(str, Enum):
+    """Supported output formats."""
+    PANDAS = "pandas"
+    JSON = "json"
+    DICT = "dict"
+    CSV = "csv"
+    STRUCTURED = "structured"  # Uses Pydantic models
+
+
+class SchemaMetadata(BaseModel):
+    """Metadata for a database schema."""
+    schema_name: str
+    tables: List[Dict[str, Any]]
+    views: List[Dict[str, Any]]
+    functions: List[Dict[str, Any]]
+    procedures: List[Dict[str, Any]]
+    indexes: List[Dict[str, Any]]
+    constraints: List[Dict[str, Any]]
+    last_updated: datetime
+    database_flavor: DatabaseFlavor
+
+
+class QueryValidationResult(BaseModel):
+    """Result of query validation."""
+    is_valid: bool
+    query_type: Optional[QueryType]
+    affected_tables: List[str]
+    estimated_cost: Optional[float]
+    warnings: List[str]
+    errors: List[str]
+    security_checks: Dict[str, bool]
+
+
+class DatabaseToolArgs(AbstractToolArgsSchema):
+    """Arguments for the unified database tool."""
+
+    # Query specification
+    natural_language_query: Optional[str] = Field(
+        None, description="Natural language description of what you want to query"
+    )
+    sql_query: Optional[str] = Field(
+        None, description="Direct SQL query to execute"
+    )
+
+    # Database connection
+    database_flavor: DatabaseFlavor = Field(
+        DatabaseFlavor.POSTGRESQL, description="Type of database to connect to"
+    )
+    connection_params: Optional[Dict[str, Any]] = Field(
+        None, description="Database connection parameters"
+    )
+    schema_names: List[str] = Field(
+        default=["public"], description="Schema names to work with"
+    )
+
+    # Operation modes
+    operation: Literal[
+        "schema_extract", "query_generate", "query_validate",
+        "query_execute", "full_pipeline", "explain_query"
+    ] = Field(
+        "full_pipeline", description="What operation to perform"
+    )
+
+    # Query options
+    max_rows: int = Field(1000, description="Maximum rows to return")
+    timeout_seconds: int = Field(300, description="Query timeout")
+    dry_run: bool = Field(False, description="Validate without executing")
+
+    # Output options
+    output_format: OutputFormat = Field(
+        OutputFormat.PANDAS, description="Format for query results"
+    )
+    structured_output_schema: Optional[Dict[str, Any]] = Field(
+        None, description="Pydantic schema for structured outputs"
+    )
+
+    # Knowledge base options
+    update_knowledge_base: bool = Field(
+        True, description="Whether to update schema knowledge base"
+    )
+    cache_duration_hours: int = Field(
+        24, description="How long to cache schema metadata"
+    )
+
+    @model_validator(mode='after')
+    def validate_query_input(self) -> 'DatabaseToolArgs':
+        # Ensure at least one query type is provided for query operations
+        if self.operation in ['query_generate', 'query_execute', 'full_pipeline', 'explain_query']:
+            if not self.natural_language_query and not self.sql_query:
+                raise ValueError("Either natural_language_query or sql_query must be provided")
+        return self
+
+
+class DatabaseTool(AbstractTool):
+    """
+    Unified Database Tool that handles the complete database interaction pipeline:
+
+    1. Schema Discovery: Extract and cache table schemas from any supported database
+    2. Knowledge Base Building: Store schema metadata in vector store for RAG
+    3. Query Generation: Convert natural language to database-specific queries
+    4. Query Validation: Syntax checking, security validation, cost estimation
+    5. Query Execution: Safe execution with proper error handling
+    6. Structured Output: Format results according to specified schemas
+
+    This tool consolidates the functionality of SchemaTool, DatabaseQueryTool,
+    and SQLAgent into a single, cohesive interface.
+    """
+
+    name = "database_tool"
+    description = """Unified database tool for schema discovery, query generation,
+                    validation, and execution across multiple database types"""
+    args_schema = DatabaseToolArgs
+
+    def __init__(
+        self,
+        knowledge_store: Optional[AbstractStore] = None,
+        default_connection_params: Optional[Dict[DatabaseFlavor, Dict]] = None,
+        enable_query_caching: bool = True,
+        llm: Optional[Union[AbstractClient, str]] = None,
+        **kwargs
+    ):
+        """
+        Initialize the unified database tool.
+
+        Args:
+            knowledge_store: Vector store for schema metadata and RAG
+            default_connection_params: Default connection parameters per database type
+            enable_query_caching: Whether to cache query results
+            llm: LLM to use for query generation and validation
+        """
+        super().__init__(**kwargs)
+
+        self.knowledge_store = knowledge_store
+        self.default_connection_params = default_connection_params or {}
+        self.enable_query_caching = enable_query_caching
+
+        # Initialize LLM
+        if isinstance(llm, str):
+            self.llm = LLMFactory.create(llm)
+        else:
+            self.llm = llm
+
+        # Cache for schema metadata and database connections
+        self._schema_cache: Dict[str, Tuple[SchemaMetadata, datetime]] = {}
+        self._connection_cache: Dict[str, AsyncDB] = {}
+
+        # Database-specific query generators and validators
+        self._query_generators = {}
+        self._query_validators = {}
+
+        self._setup_database_handlers()
+
+    def _setup_database_handlers(self):
+        """Initialize database-specific handlers for different flavors."""
+        # This would be expanded to include handlers for each database type
+        self._query_generators = {
+            DatabaseFlavor.POSTGRESQL: self._generate_postgresql_query,
+            DatabaseFlavor.MYSQL: self._generate_mysql_query,
+            DatabaseFlavor.BIGQUERY: self._generate_bigquery_query,
+            # Add more database-specific generators...
+        }
+
+        self._query_validators = {
+            DatabaseFlavor.POSTGRESQL: self._validate_postgresql_query,
+            DatabaseFlavor.MYSQL: self._validate_mysql_query,
+            DatabaseFlavor.BIGQUERY: self._validate_bigquery_query,
+            # Add more database-specific validators...
+        }
+
+    def _clean_sql(self, sql_query: str) -> str:
+        """Clean SQL query from markdown formatting."""
+        if not sql_query:
+            return ""
+        # Remove markdown code blocks
+        clean_query = re.sub(r'```\w*\n?', '', sql_query)
+        clean_query = clean_query.replace('```', '')
+        return clean_query.strip()
+
+    async def _execute(
+        self,
+        natural_language_query: Optional[str] = None,
+        sql_query: Optional[str] = None,
+        database_flavor: DatabaseFlavor = DatabaseFlavor.POSTGRESQL,
+        connection_params: Optional[Dict[str, Any]] = None,
+        schema_names: List[str] = ["public"],
+        operation: str = "full_pipeline",
+        max_rows: int = 1000,
+        timeout_seconds: int = 300,
+        dry_run: bool = False,
+        output_format: OutputFormat = OutputFormat.PANDAS,
+        structured_output_schema: Optional[Dict[str, Any]] = None,
+        update_knowledge_base: bool = True,
+        cache_duration_hours: int = 24,
+        **kwargs
+    ) -> ToolResult:
+        """
+        Execute the unified database tool pipeline.
+
+        The method routes to different sub-operations based on the operation parameter,
+        or executes the full pipeline for complete query processing.
+        """
+        try:
+            # Fallback to default connection parameters if not provided
+            if connection_params is None:
+                connection_params = self.default_connection_params.get(database_flavor)
+
+            if sql_query:
+                sql_query = self._clean_sql(sql_query)
+
+            # Route to specific operations
+            if operation == "schema_extract":
+                return await self._extract_schema_operation(
+                    database_flavor, connection_params, schema_names,
+                    update_knowledge_base, cache_duration_hours
+                )
+            if operation == "query_generate":
+                return await self._query_generation_operation(
+                    natural_language_query, database_flavor, connection_params, schema_names
+                )
+            if operation == "query_validate":
+                return await self._query_validation_operation(
+                    sql_query or natural_language_query, database_flavor, connection_params
+                )
+            if operation == "query_execute":
+                return await self._query_execution_operation(
+                    sql_query, database_flavor, connection_params,
+                    max_rows, timeout_seconds, output_format, structured_output_schema
+                )
+            if operation == "full_pipeline":
+                return await self._full_pipeline_operation(
+                    natural_language_query, sql_query, database_flavor, connection_params,
+                    schema_names, max_rows, timeout_seconds, dry_run,
+                    output_format, structured_output_schema, update_knowledge_base, cache_duration_hours
+                )
+            if operation == "explain_query":
+                return await self._explain_query_operation(
+                    sql_query or natural_language_query, database_flavor, connection_params
+                )
+            else:
+                raise ValueError(f"Unknown operation: {operation}")
+
+        except Exception as e:
+            return ToolResult(
+                status="error",
+                result=None,
+                error=f"Database tool execution failed: {str(e)}",
+                metadata={
+                    "operation": operation,
+                    "database_flavor": database_flavor.value,
+                    "timestamp": datetime.now(timezone.utc).isoformat()
+                }
+            )
+
+    async def _full_pipeline_operation(
+        self,
+        natural_language_query: Optional[str],
+        sql_query: Optional[str],
+        database_flavor: DatabaseFlavor,
+        connection_params: Optional[Dict[str, Any]],
+        schema_names: List[str],
+        max_rows: int,
+        timeout_seconds: int,
+        dry_run: bool,
+        output_format: OutputFormat,
+        structured_output_schema: Optional[Dict[str, Any]],
+        update_knowledge_base: bool,
+        cache_duration_hours: int
+    ) -> ToolResult:
+        """
+        Execute the complete database interaction pipeline.
+
+        This is the main orchestrator method that combines all functionality:
+        schema extraction, knowledge base updates, query generation, validation, and execution.
+        """
+        pipeline_results = {
+            "schema_extraction": None,
+            "query_generation": None,
+            "query_validation": None,
+            "query_execution": None,
+            "knowledge_base_update": None
+        }
+
+        try:
+            # Step 1: Extract and cache schema metadata
+            self.logger.info(f"Step 1: Extracting schema for {database_flavor.value}")
+            schema_result = await self._extract_schema_operation(
+                database_flavor, connection_params, schema_names,
+                update_knowledge_base, cache_duration_hours
+            )
+            pipeline_results["schema_extraction"] = schema_result.result
+
+            # Step 2: Generate SQL query if natural language was provided
+            generated_query = sql_query
+            if natural_language_query:
+                self.logger.info("Step 2: Generating SQL from natural language")
+                query_result = await self._query_generation_operation(
+                    natural_language_query, database_flavor, connection_params, schema_names
+                )
+                pipeline_results["query_generation"] = query_result.result
+                generated_query = query_result.result.get("sql_query")
+
+            if not generated_query:
+                raise ValueError("No valid SQL query to execute")
+
+            # Step 3: Validate the query
+            self.logger.info("Step 3: Validating SQL query")
+            validation_result = await self._query_validation_operation(
+                generated_query, database_flavor, connection_params
+            )
+            pipeline_results["query_validation"] = validation_result.result
+
+            if not validation_result.result["is_valid"]:
+                if dry_run:
+                    return ToolResult(
+                        status="success",
+                        result={
+                            "pipeline_results": pipeline_results,
+                            "dry_run": True,
+                            "query_valid": False
+                        },
+                        metadata={"operation": "full_pipeline", "dry_run": True}
+                    )
+                else:
+                    raise ValueError(f"Query validation failed: {validation_result.result['errors']}")
+
+            # Step 4: Execute the query (unless dry run)
+            if not dry_run:
+                self.logger.info("Step 4: Executing validated query")
+                execution_result = await self._query_execution_operation(
+                    generated_query, database_flavor, connection_params,
+                    max_rows, timeout_seconds, output_format, structured_output_schema
+                )
+                pipeline_results["query_execution"] = execution_result.result
+
+            # Success! Return comprehensive results
+            return ToolResult(
+                status="success",
+                result={
+                    "pipeline_results": pipeline_results,
+                    "final_query": generated_query,
+                    "dry_run": dry_run,
+                    "execution_summary": {
+                        "rows_returned": len(pipeline_results["query_execution"]["data"]) if not dry_run and pipeline_results["query_execution"] else 0,
+                        "execution_time_seconds": pipeline_results["query_execution"]["execution_time"] if not dry_run and pipeline_results["query_execution"] else None,
+                        "output_format": output_format.value
+                    }
+                },
+                metadata={
+                    "operation": "full_pipeline",
+                    "database_flavor": database_flavor.value,
+                    "schema_count": len(schema_names),
+                    "natural_language_input": natural_language_query is not None,
+                    "timestamp": datetime.utcnow().isoformat()
+                }
+            )
+
+        except Exception as e:
+            return ToolResult(
+                status="error",
+                result={"pipeline_results": pipeline_results},
+                error=f"Pipeline failed at step: {str(e)}",
+                metadata={"operation": "full_pipeline", "partial_results": True}
+            )
+
+    async def _extract_schema_operation(
+        self,
+        database_flavor: DatabaseFlavor,
+        connection_params: Optional[Dict[str, Any]],
+        schema_names: List[str],
+        update_knowledge_base: bool,
+        cache_duration_hours: int
+    ) -> ToolResult:
+        """Extract database schema metadata and optionally update knowledge base."""
+        try:
+            # Check cache first
+            cache_key = self._generate_schema_cache_key(database_flavor, connection_params, schema_names)
+            cached_schema, cache_time = self._schema_cache.get(cache_key, (None, None))
+
+            if cached_schema and cache_time:
+                cache_age = datetime.utcnow() - cache_time
+                if cache_age < timedelta(hours=cache_duration_hours):
+                    self.logger.info(f"Using cached schema metadata (age: {cache_age})")
+                    return ToolResult(
+                        status="success",
+                        result=cached_schema.dict(),
+                        metadata={"source": "cache", "cache_age_hours": cache_age.total_seconds() / 3600}
+                    )
+
+            # Extract fresh schema metadata
+            db_connection = await self._get_database_connection(database_flavor, connection_params)
+            schema_metadata = await self._extract_database_schema(db_connection, database_flavor, schema_names)
+
+            # Cache the results
+            self._schema_cache[cache_key] = (schema_metadata, datetime.utcnow())
+
+            # Update knowledge base if requested
+            if update_knowledge_base and self.knowledge_store:
+                await self._update_schema_knowledge_base(schema_metadata)
+
+            return ToolResult(
+                status="success",
+                result=schema_metadata.dict(),
+                metadata={
+                    "source": "database",
+                    "schema_count": len(schema_names),
+                    "table_count": len(schema_metadata.tables),
+                    "view_count": len(schema_metadata.views),
+                    "knowledge_base_updated": update_knowledge_base and self.knowledge_store is not None
+                }
+            )
+
+        except Exception as e:
+            return ToolResult(
+                status="error",
+                result=None,
+                error=f"Schema extraction failed: {str(e)}",
+                metadata={"operation": "schema_extract"}
+            )
+
+    # Additional helper methods would continue here...
+    # Including _query_generation_operation, _query_validation_operation,
+    # _query_execution_operation, and all the database-specific implementations
+
+    def _generate_schema_cache_key(
+        self,
+        database_flavor: DatabaseFlavor,
+        connection_params: Optional[Dict[str, Any]],
+        schema_names: List[str]
+    ) -> str:
+        """Generate a unique cache key for schema metadata."""
+        key_data = {
+            "flavor": database_flavor.value,
+            "params": connection_params or {},
+            "schemas": sorted(schema_names)
+        }
+        return hashlib.md5(json.dumps(key_data, sort_keys=True).encode()).hexdigest()
+
+    async def _get_database_connection(
+        self,
+        database_flavor: DatabaseFlavor,
+        connection_params: Optional[Dict[str, Any]]
+    ) -> AsyncDB:
+        """Get or create a database connection using AsyncDB."""
+        """Get or create a database connection using AsyncDB."""
+        # Normalize connection parameters
+        params = connection_params.copy() if connection_params else {}
+
+        # Common mapping: username -> user (used by asyncpg and others)
+        if 'username' in params and 'user' not in params:
+            params['user'] = params.pop('username')
+
+        driver_map = {
+            DatabaseFlavor.POSTGRESQL: 'pg',
+            DatabaseFlavor.MYSQL: 'mysql',
+            DatabaseFlavor.SQLITE: 'sqlite',
+        }
+        driver = driver_map.get(database_flavor, database_flavor.value)
+        return AsyncDB(driver, params=params)
+
+    async def _extract_database_schema(
+        self,
+        db_connection: AsyncDB,
+        database_flavor: DatabaseFlavor,
+        schema_names: List[str]
+    ) -> SchemaMetadata:
+        """Extract comprehensive schema metadata from the database."""
+        """Extract comprehensive schema metadata from the database."""
+        if database_flavor == DatabaseFlavor.POSTGRESQL:
+            return await self._extract_postgresql_schema(db_connection, schema_names)
+
+        raise NotImplementedError(f"Schema extraction not implemented for {database_flavor}")
+
+    async def _extract_postgresql_schema(
+        self,
+        db: AsyncDB,
+        schema_names: List[str]
+    ) -> SchemaMetadata:
+        """Extract schema for PostgreSQL."""
+        tables_data = []
+        async with await db.connection() as conn:
+            schemas_list = ", ".join([f"'{s}'" for s in schema_names])
+            if not schemas_list:
+                schemas_list = "'public'" # Default
+
+            query = f"""
+                SELECT t.table_schema, t.table_name, c.column_name, c.data_type
+                FROM information_schema.tables t
+                JOIN information_schema.columns c
+                  ON t.table_schema = c.table_schema AND t.table_name = c.table_name
+                WHERE t.table_schema IN ({schemas_list})
+                ORDER BY t.table_schema, t.table_name, c.ordinal_position
+            """
+            try:
+                rows = await conn.fetch(query) # Using fetch if available, or query
+            except Exception:
+                # Fallback to query if fetch not available on conn wrapper
+                rows = await conn.query(query)
+
+            # Check if rows is a list of lists (result set wrapper)
+            if rows and isinstance(rows, list) and len(rows) > 0 and isinstance(rows[0], list):
+                rows = rows[0]
+
+            # Process rows
+            grouped = {}
+            for row in rows:
+                # Handle possible dict or object access
+                # asyncpg.Record supports .get() and ['key']
+                if hasattr(row, 'get'):
+                    s_name = row.get('table_schema')
+                    t_name = row.get('table_name')
+                    c_name = row.get('column_name')
+                    d_type = row.get('data_type')
+                elif isinstance(row, (list, tuple)) and len(row) >= 4:
+                    s_name = row[0]
+                    t_name = row[1]
+                    c_name = row[2]
+                    d_type = row[3]
+                else:
+                    # Attempt dict access as fallback
+                    try:
+                        s_name = row['table_schema']
+                        t_name = row['table_name']
+                        c_name = row['column_name']
+                        d_type = row['data_type']
+                    except (TypeError, KeyError, IndexError):
+                        continue # Skip invalid rows
+
+                k = (s_name, t_name)
+                if k not in grouped:
+                    grouped[k] = {
+                        "schema": s_name,
+                        "name": t_name,
+                        "columns": []
+                    }
+                grouped[k]["columns"].append({"name": c_name, "type": d_type})
+
+            tables_data = list(grouped.values())
+
+        return SchemaMetadata(
+            schema_name=",".join(schema_names),
+            tables=tables_data,
+            views=[],
+            functions=[],
+            procedures=[],
+            indexes=[],
+            constraints=[],
+            last_updated=datetime.utcnow(),
+            database_flavor=DatabaseFlavor.POSTGRESQL
+        )
+
+    async def _query_generation_operation(
+        self,
+        natural_language_query: str,
+        database_flavor: DatabaseFlavor,
+        connection_params: Optional[Dict[str, Any]],
+        schema_names: List[str]
+    ) -> ToolResult:
+        """Generate SQL query from natural language using schema context."""
+        try:
+            # Get schema context for query generation
+            schema_key = self._generate_schema_cache_key(database_flavor, connection_params, schema_names)
+            cached_schema, _ = self._schema_cache.get(schema_key, (None, None))
+
+            if not cached_schema:
+                # If no cached schema, extract it first
+                schema_result = await self._extract_schema_operation(
+                    database_flavor, connection_params, schema_names, False, 24
+                )
+                if schema_result.status != "success" or not schema_result.result:
+                    raise ValueError(f"Schema extraction failed: {schema_result.error or 'No result returned'}")
+
+                cached_schema = SchemaMetadata(**schema_result.result)
+
+            # Use database-specific query generator
+            generator = self._query_generators.get(database_flavor)
+            if not generator:
+                raise ValueError(f"No query generator available for {database_flavor.value}")
+
+            # Build rich context for LLM query generation
+            schema_context = self._build_schema_context_for_llm(cached_schema, natural_language_query)
+
+            # Generate the SQL query
+            generated_sql = await generator(natural_language_query, schema_context)
+            generated_sql = self._clean_sql(generated_sql)
+
+            return ToolResult(
+                status="success",
+                result={
+                    "natural_language_query": natural_language_query,
+                    "sql_query": generated_sql,
+                    "database_flavor": database_flavor.value,
+                    "schema_context_used": len(schema_context.get("relevant_tables", [])),
+                    "generation_timestamp": datetime.utcnow().isoformat()
+                },
+                metadata={
+                    "operation": "query_generation",
+                    "has_schema_context": bool(schema_context)
+                }
+            )
+
+        except Exception as e:
+            return ToolResult(
+                status="error",
+                result=None,
+                error=f"Query generation failed: {str(e)}",
+                metadata={"operation": "query_generation"}
+            )
+
+    async def _query_validation_operation(
+        self,
+        sql_query: str,
+        database_flavor: DatabaseFlavor,
+        connection_params: Optional[Dict[str, Any]]
+    ) -> ToolResult:
+        """Validate SQL query for syntax, security, and performance."""
+        try:
+            validator = self._query_validators.get(database_flavor)
+            if not validator:
+                raise ValueError(f"No query validator available for {database_flavor.value}")
+
+            validation_result = await validator(sql_query)
+
+            return ToolResult(
+                status="success" if validation_result.is_valid else "warning",
+                result=validation_result.dict(),
+                metadata={
+                    "operation": "query_validation",
+                    "query_type": validation_result.query_type.value if validation_result.query_type else None
+                }
+            )
+
+        except Exception as e:
+            return ToolResult(
+                status="error",
+                result=None,
+                error=f"Query validation failed: {str(e)}",
+                metadata={"operation": "query_validation"}
+            )
+
+    async def _query_execution_operation(
+        self,
+        sql_query: str,
+        database_flavor: DatabaseFlavor,
+        connection_params: Optional[Dict[str, Any]],
+        max_rows: int,
+        timeout_seconds: int,
+        output_format: OutputFormat,
+        structured_output_schema: Optional[Dict[str, Any]]
+    ) -> ToolResult:
+        """Execute SQL query and format results according to specifications."""
+        try:
+            db_connection = await self._get_database_connection(database_flavor, connection_params)
+
+            # Execute query with timeout and row limit
+            start_time = datetime.utcnow()
+
+            # This integrates your existing DatabaseQueryTool logic
+            raw_results = await self._execute_query_with_asyncdb(
+                db_connection, sql_query, max_rows, timeout_seconds
+            )
+
+            execution_time = (datetime.utcnow() - start_time).total_seconds()
+
+            # Format results according to specified output format
+            formatted_results = await self._format_query_results(
+                raw_results, output_format, structured_output_schema
+            )
+
+            return ToolResult(
+                status="success",
+                result={
+                    "data": formatted_results,
+                    "row_count": len(raw_results) if isinstance(raw_results, list) else None,
+                    "execution_time": execution_time,
+                    "output_format": output_format.value,
+                    "query": sql_query
+                },
+                metadata={
+                    "operation": "query_execution",
+                    "database_flavor": database_flavor.value,
+                    "rows_returned": len(raw_results) if isinstance(raw_results, list) else 0
+                }
+            )
+
+        except Exception as e:
+            return ToolResult(
+                status="error",
+                result=None,
+                error=f"Query execution failed: {str(e)}",
+                metadata={"operation": "query_execution", "query": sql_query}
+            )
+
+    async def _explain_query_operation(
+        self,
+        sql_query: str,
+        database_flavor: DatabaseFlavor,
+        connection_params: Optional[Dict[str, Any]]
+    ) -> ToolResult:
+        """
+        Explain query execution plan and provide LLM-based optimizations.
+        """
+        if not sql_query:
+             return ToolResult(
+                status="error",
+                result=None,
+                error="No SQL query provided for explanation",
+                metadata={"operation": "explain_query"}
+            )
+
+        try:
+            db_connection = await self._get_database_connection(database_flavor, connection_params)
+
+            # Determine appropriate EXPLAIN command
+            explain_cmd = f"EXPLAIN ANALYZE {sql_query}"
+            if database_flavor == DatabaseFlavor.MYSQL:
+                 # MySQL 8.0.18+ supports EXPLAIN ANALYZE, otherwise fallback to EXPLAIN
+                 # For safety/compatibility we might start with EXPLAIN if ANALYZE fails or just try
+                 explain_cmd = f"EXPLAIN ANALYZE {sql_query}"
+            elif database_flavor == DatabaseFlavor.BIGQUERY:
+                 # BigQuery doesn't support EXPLAIN ANALYZE syntax directly in this way usually
+                 # It returns stats in job metadata.
+                 # But we can try to use Dry Run or similar.
+                 # For now, let's assume standard SQL syntax applies or let execution fail and fallback
+                 pass
+
+            # Execute explanation
+            try:
+                raw_plan = await self._execute_query_with_asyncdb(
+                    db_connection, explain_cmd, max_rows=0, timeout_seconds=30
+                )
+            except Exception as e:
+                # Fallback to simple EXPLAIN if ANALYZE fails (e.g. not supported or timeouts)
+                self.logger.warning(f"EXPLAIN ANALYZE failed, falling back to EXPLAIN: {e}")
+                explain_cmd = f"EXPLAIN {sql_query}"
+                raw_plan = await self._execute_query_with_asyncdb(
+                    db_connection, explain_cmd, max_rows=0, timeout_seconds=30
+                )
+
+            # Format plan into string
+            plan_text = ""
+            if isinstance(raw_plan, list):
+                # Flatten the list of rows/dicts
+                for row in raw_plan:
+                    if isinstance(row, dict):
+                        # Usually the first column contains the plan output
+                        plan_text += list(row.values())[0] + "\n"
+                    elif isinstance(row, (list, tuple)):
+                        plan_text += str(row[0]) + "\n"
+                    else:
+                        plan_text += str(row) + "\n"
+            else:
+                plan_text = str(raw_plan)
+
+            # Ask LLM to explain and optimize
+            llm_explanation = "No LLM configured for explanation."
+            if self.llm:
+                prompt = (
+                    f"You are a database performance expert. Analyze the following query plan for a {database_flavor.value} database.\n"
+                    f"Query:\n```sql\n{sql_query}\n```\n\n"
+                    f"Execution Plan:\n```\n{plan_text}\n```\n\n"
+                    "Please provide:\n"
+                    "1. A human-readable explanation of how the query is executed.\n"
+                    "2. Performance bottlenecks identified in the plan.\n"
+                    "3. Concrete suggestions for indexes or query rewrites to improve performance.\n"
+                    "4. Rating of current query efficiency (1-10)."
+                )
+
+                response = await self.llm.ask(prompt)
+                if isinstance(response, AIMessage):
+                    llm_explanation = str(response.output).strip()
+                elif isinstance(response, dict) and 'content' in response:
+                    llm_explanation = str(response['content']).strip()
+                else:
+                    llm_explanation = str(response).strip()
+
+            return ToolResult(
+                status="success",
+                result={
+                    "query": sql_query,
+                    "plan": plan_text,
+                    "analysis": llm_explanation,
+                    "database_flavor": database_flavor.value
+                },
+                metadata={
+                    "operation": "explain_query",
+                    "command_used": explain_cmd
+                }
+            )
+
+        except Exception as e:
+            return ToolResult(
+                status="error",
+                result=None,
+                error=f"Query explanation failed: {str(e)}",
+                metadata={"operation": "explain_query"}
+            )
+
+    def _build_schema_context_for_llm(
+        self,
+        schema_metadata: SchemaMetadata,
+        natural_language_query: str
+    ) -> Dict[str, Any]:
+        """
+        Build rich schema context for LLM query generation.
+
+        This is a critical method that determines query generation quality.
+        It intelligently selects relevant schema elements based on the natural language query.
+        """
+        # Use vector similarity or keyword matching to find relevant tables
+        relevant_tables = self._find_relevant_tables(schema_metadata, natural_language_query)
+
+        # Build comprehensive context including relationships, constraints, and sample data
+        context = {
+            "database_flavor": schema_metadata.database_flavor.value,
+            "schema_name": schema_metadata.schema_name,
+            "relevant_tables": relevant_tables,
+            "table_relationships": self._extract_table_relationships(schema_metadata, relevant_tables),
+            "common_patterns": self._get_query_patterns_for_tables(relevant_tables),
+            "data_types_guide": self._get_data_type_guide(schema_metadata.database_flavor)
+        }
+
+        return context
+
+    async def _execute_query_with_asyncdb(
+        self,
+        db_connection: AsyncDB,
+        sql_query: str,
+        max_rows: int,
+        timeout_seconds: int
+    ) -> Any:
+        """Execute query using AsyncDB with proper error handling and limits."""
+        # This integrates your existing DatabaseQueryTool execution logic
+        # but with enhanced error handling and result limiting
+
+        try:
+            # Add LIMIT clause if not present and max_rows is specified
+            if max_rows > 0 and "LIMIT" not in sql_query.upper():
+                sql_query = f"{sql_query.rstrip(';')} LIMIT {max_rows};"
+
+            # Execute with timeout using asyncio
+            async with await db_connection.connection() as conn:
+                return await asyncio.wait_for(
+                    conn.fetchall(sql_query),
+                    timeout=timeout_seconds
+                )
+
+        except asyncio.TimeoutError as e:
+            raise TimeoutError(
+                f"Query execution timed out after {timeout_seconds} seconds"
+            ) from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Database execution error: {str(e)}"
+            ) from e
+
+    async def _format_query_results(
+        self,
+        raw_results: Any,
+        output_format: OutputFormat,
+        structured_output_schema: Optional[Dict[str, Any]]
+    ) -> Any:
+        """Format query results according to specified output format."""
+        if output_format == OutputFormat.PANDAS:
+            return pd.DataFrame(raw_results) if raw_results else pd.DataFrame()
+        elif output_format == OutputFormat.JSON:
+            return json.dumps(raw_results, default=str, indent=2)
+        elif output_format == OutputFormat.DICT:
+            return raw_results
+        elif output_format == OutputFormat.CSV:
+            df = pd.DataFrame(raw_results) if raw_results else pd.DataFrame()
+            return df.to_csv(index=False)
+        elif output_format == OutputFormat.STRUCTURED and structured_output_schema:
+            # Convert results to Pydantic models based on provided schema
+            return self._convert_to_structured_output(raw_results, structured_output_schema)
+        else:
+            return raw_results
+
+    # Database-specific implementations (these would replace your current separate tools)
+    async def _generate_postgresql_query(self, natural_language: str, schema_context: Dict) -> str:
+        """
+        Generate PostgreSQL-specific SQL from natural language.
+
+        This method would integrate your existing SQLAgent logic but with enhanced
+        schema context and PostgreSQL-specific optimizations.
+        """
+        # Build prompt with rich schema context
+        prompt = self._build_query_generation_prompt(
+            natural_language, schema_context, "postgresql"
+        )
+
+        # Use your existing LLM client to generate the query
+        # This would integrate with your AI-Parrot LLM clients
+        return await self._call_llm_for_query_generation(prompt)
+
+    async def _validate_postgresql_query(self, query: str) -> QueryValidationResult:
+        """
+        Validate PostgreSQL query for syntax, security, and performance.
+
+        This provides the validation layer that was missing from your current SQLAgent.
+        """
+        validation_result = QueryValidationResult(
+            is_valid=True,
+            query_type=None,
+            affected_tables=[],
+            estimated_cost=None,
+            warnings=[],
+            errors=[],
+            security_checks={}
+        )
+
+        try:
+            # Parse query to determine type and affected tables
+            query_upper = query.strip().upper()
+            if query_upper.startswith('SELECT'):
+                validation_result.query_type = QueryType.SELECT
+            elif query_upper.startswith('INSERT'):
+                validation_result.query_type = QueryType.INSERT
+            # ... other query types
+
+            # Security checks
+            validation_result.security_checks = {
+                "no_sql_injection_patterns": self._check_sql_injection_patterns(query),
+                "no_dangerous_operations": self._check_dangerous_operations(query),
+                "proper_quoting": self._check_proper_quoting(query)
+            }
+
+            # Syntax validation (could use sqlparse or connect to database for EXPLAIN)
+            syntax_valid = await self._validate_syntax_postgresql(query)
+            if not syntax_valid:
+                validation_result.is_valid = False
+                validation_result.errors.append("Invalid SQL syntax")
+
+            # Performance warnings
+            if "SELECT *" in query_upper:
+                validation_result.warnings.append("Consider specifying explicit columns instead of SELECT *")
+
+            return validation_result
+
+        except Exception as e:
+            validation_result.is_valid = False
+            validation_result.errors.append(f"Validation error: {str(e)}")
+            return validation_result
+
+    def _check_dangerous_operations(self, query: str) -> bool:
+        """
+        Check if query contains dangerous operations that should be blocked.
+
+        Returns:
+            True if query is SAFE (no dangerous operations)
+            False if dangerous operations detected
+        """
+        query_upper = query.upper()
+
+        dangerous_patterns = [
+            # DDL operations
+            r'\bDROP\s+(TABLE|DATABASE|SCHEMA|INDEX|VIEW|PROCEDURE|FUNCTION)\b',
+            r'\bTRUNCATE\s+TABLE\b',
+            r'\bALTER\s+(TABLE|DATABASE|SCHEMA)\s+.*\s+DROP\b',
+            # DML without WHERE
+            r'\bDELETE\s+FROM\s+\w+\s*;?\s*$',
+            # Admin commands
+            r'\bGRANT\b',
+            r'\bREVOKE\b',
+            r'\bCREATE\s+USER\b',
+            r'\bDROP\s+USER\b',
+            r'\bALTER\s+USER\b',
+            # Command execution (PostgreSQL)
+            r'\bCOPY\s+.*\s+TO\s+PROGRAM\b',
+            # SQL Server
+            r'\bEXEC\s*\(',
+            r'\bXP_CMDSHELL\b',
+            # MySQL file operations
+            r'\bLOAD_FILE\b',
+            r'\bINTO\s+OUTFILE\b',
+            r'\bINTO\s+DUMPFILE\b',
+        ]
+
+        for pattern in dangerous_patterns:
+            if re.search(pattern, query_upper, re.IGNORECASE | re.DOTALL):
+                return False
+
+        # Check DELETE/UPDATE without WHERE
+        if re.search(r'\bDELETE\s+FROM\s+\w+\s*$', query_upper):
+            return False
+
+        update_match = re.search(r'\bUPDATE\s+\w+\s+SET\s+', query_upper)
+        if update_match and 'WHERE' not in query_upper:
+            return False
+
+        return True
+
+    def _check_sql_injection_patterns(self, query: str) -> bool:
+        """
+        Check for common SQL injection patterns.
+
+        Returns:
+            True if no injection patterns found (SAFE)
+            False if potential injection detected
+        """
+        injection_patterns = [
+            # Union/Boolean-based
+            r"'\s*(OR|AND)\s+['\"0-9]",
+            r"'\s*OR\s+1\s*=\s*1",
+            r"'\s*OR\s+'[^']*'\s*=\s*'[^']*'",
+            # Comment-based
+            r";\s*--",
+            r";\s*/\*",
+            r"--\s*$",
+            # Stacked queries
+            r"'\s*;\s*(DROP|DELETE|UPDATE|INSERT|EXEC)\b",
+            r";\s*(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER)\b",
+            # UNION injection
+            r"\bUNION\s+(ALL\s+)?SELECT\b.*\bFROM\b",
+            # Time-based
+            r"\bSLEEP\s*\(",
+            r"\bWAITFOR\s+DELAY\b",
+            r"\bBENCHMARK\s*\(",
+            r"\bPG_SLEEP\s*\(",
+            # Encoding attempts
+            r"0x[0-9a-fA-F]+",
+            r"\bCHAR\s*\(\s*\d+\s*\)",
+        ]
+
+        for pattern in injection_patterns:
+            if re.search(pattern, query, re.IGNORECASE):
+                return False
+
+        return True
+
+    def _check_proper_quoting(self, query: str) -> bool:
+        """
+        Check if string literals are properly quoted.
+
+        Returns:
+            True if quoting appears proper (SAFE)
+            False if improper quoting detected
+        """
+        # Check for unbalanced quotes
+        single_quotes = query.count("'") - query.count("\\'") - query.count("''")
+        double_quotes = query.count('"') - query.count('\\"') - query.count('""')
+
+        if single_quotes % 2 != 0 or double_quotes % 2 != 0:
+            return False
+
+        # Check for dangerous patterns after string literals
+        dangerous_patterns = [
+            r"'\s*\)\s*(OR|AND|UNION)\b",
+            r"'\s*;\s*\w+",
+        ]
+
+        for pattern in dangerous_patterns:
+            if re.search(pattern, query, re.IGNORECASE):
+                return False
+
+        return True
+
+    async def _validate_syntax_postgresql(self, query: str) -> bool:
+        """Validate PostgreSQL query syntax using pattern matching."""
+        try:
+            query_stripped = query.strip().rstrip(';')
+            query_upper = query_stripped.upper()
+
+            valid_starts = [
+                'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'WITH',
+                'CREATE', 'ALTER', 'DROP', 'TRUNCATE',
+                'BEGIN', 'COMMIT', 'ROLLBACK', 'SAVEPOINT',
+                'EXPLAIN', 'ANALYZE', 'VACUUM', 'REINDEX',
+                'GRANT', 'REVOKE', 'SET', 'SHOW', 'RESET',
+                'COPY', 'CALL', 'DO', 'LOCK',
+            ]
+
+            first_word = query_upper.split()[0] if query_stripped else ''
+            if first_word not in valid_starts:
+                return False
+
+            # Validate statement structures
+            if query_upper.startswith('INSERT') and 'INTO' not in query_upper:
+                return False
+            if query_upper.startswith('UPDATE') and 'SET' not in query_upper:
+                return False
+            if query_upper.startswith('DELETE') and 'FROM' not in query_upper:
+                return False
+
+            # Check balance
+            if query.count('(') != query.count(')'):
+                return False
+            if query.count('[') != query.count(']'):
+                return False
+
+            return True
+        except Exception:
+            return False
+
+    # =========================================================================
+    # MYSQL METHODS
+    # =========================================================================
+
+    async def _generate_mysql_query(
+        self,
+        natural_language: str,
+        schema_context: Dict[str, Any]
+    ) -> str:
+        """Generate MySQL-specific SQL from natural language."""
+        prompt = self._build_query_generation_prompt(
+            natural_language, schema_context, "mysql"
+        )
+
+        mysql_instructions = """
+MySQL-Specific Rules:
+1. Use backticks (`) for identifier quoting
+2. Use LIMIT for row limiting
+3. Use IFNULL() instead of COALESCE() for two arguments
+4. Use NOW() for current timestamp
+5. Use DATE_FORMAT() for date formatting
+6. Use CONCAT() for string concatenation
+7. Boolean values are 1/0
+8. Use REGEXP for regex matching
+"""
+
+        generated_query = await self._call_llm_for_query_generation(
+            f"{prompt}\n\n{mysql_instructions}"
+        )
+
+        return self._ensure_mysql_compatibility(generated_query)
+
+    def _ensure_mysql_compatibility(self, query: str) -> str:
+        """Post-process query to ensure MySQL compatibility."""
+        result = query
+
+        # Replace double quotes with backticks for identifiers
+        result = re.sub(r'"(\w+)"(?=\s*[,.\)\s]|$)', r'`\1`', result)
+
+        # Replace COALESCE with two args to IFNULL
+        result = re.sub(
+            r'\bCOALESCE\s*\(\s*([^,]+)\s*,\s*([^,\)]+)\s*\)',
+            r'IFNULL(\1, \2)',
+            result,
+            flags=re.IGNORECASE
+        )
+
+        # Replace TRUE/FALSE with 1/0
+        result = re.sub(r'\bTRUE\b', '1', result, flags=re.IGNORECASE)
+        result = re.sub(r'\bFALSE\b', '0', result, flags=re.IGNORECASE)
+
+        return result
+
+    async def _validate_mysql_query(self, query: str) -> QueryValidationResult:
+        """Validate MySQL query for syntax, security, and performance."""
+        validation_result = QueryValidationResult(
+            is_valid=True,
+            query_type=None,
+            affected_tables=[],
+            estimated_cost=None,
+            warnings=[],
+            errors=[],
+            security_checks={}
+        )
+
+        try:
+            query_upper = query.strip().upper()
+
+            # Determine query type
+            for qt in QueryType:
+                if query_upper.startswith(qt.value):
+                    validation_result.query_type = qt
+                    break
+
+            # Extract tables
+            validation_result.affected_tables = self._extract_tables_from_query(query)
+
+            # Security checks
+            validation_result.security_checks = {
+                "no_sql_injection_patterns": self._check_sql_injection_patterns(query),
+                "no_dangerous_operations": self._check_dangerous_operations(query),
+                "proper_quoting": self._check_proper_quoting(query)
+            }
+
+            if not all(validation_result.security_checks.values()):
+                validation_result.is_valid = False
+                for check, passed in validation_result.security_checks.items():
+                    if not passed:
+                        validation_result.errors.append(f"Security check failed: {check}")
+
+            # Syntax validation
+            if not await self._validate_syntax_mysql(query):
+                validation_result.is_valid = False
+                validation_result.errors.append("Invalid MySQL syntax")
+
+            # Performance warnings
+            if "SELECT *" in query_upper:
+                validation_result.warnings.append("Consider specifying explicit columns")
+
+            if re.search(r"LIKE\s*'%", query, re.IGNORECASE):
+                validation_result.warnings.append("Leading wildcard may prevent index usage")
+
+            return validation_result
+
+        except Exception as e:
+            validation_result.is_valid = False
+            validation_result.errors.append(f"Validation error: {str(e)}")
+            return validation_result
+
+    async def _validate_syntax_mysql(self, query: str) -> bool:
+        """Validate MySQL-specific query syntax."""
+        try:
+            query_stripped = query.strip().rstrip(';')
+            query_upper = query_stripped.upper()
+
+            valid_starts = [
+                'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'REPLACE',
+                'CREATE', 'ALTER', 'DROP', 'TRUNCATE', 'RENAME',
+                'START', 'BEGIN', 'COMMIT', 'ROLLBACK', 'SAVEPOINT',
+                'SET', 'SHOW', 'DESCRIBE', 'DESC', 'EXPLAIN',
+                'GRANT', 'REVOKE', 'LOCK', 'UNLOCK', 'USE', 'WITH',
+            ]
+
+            first_word = query_upper.split()[0] if query_stripped else ''
+            if first_word not in valid_starts:
+                return False
+
+            # Validate statement structures
+            if first_word in ('INSERT', 'REPLACE') and 'INTO' not in query_upper:
+                return False
+            if first_word == 'UPDATE' and 'SET' not in query_upper:
+                return False
+
+            # Check balance
+            if query.count('(') != query.count(')'):
+                return False
+            if query.count('`') % 2 != 0:
+                return False
+
+            return True
+        except Exception:
+            return False
+
+    # =========================================================================
+    # BIGQUERY METHODS
+    # =========================================================================
+
+    async def _generate_bigquery_query(
+        self,
+        natural_language: str,
+        schema_context: Dict[str, Any]
+    ) -> str:
+        """Generate BigQuery-specific SQL from natural language."""
+        prompt = self._build_query_generation_prompt(
+            natural_language, schema_context, "bigquery"
+        )
+
+        bigquery_instructions = """
+BigQuery-Specific Rules:
+1. Use backticks for table names: `project.dataset.table`
+2. Use STRUCT<> and ARRAY<> for complex types
+3. Use UNNEST() to flatten arrays
+4. Use SAFE_DIVIDE() for division with potential zeros
+5. Use FORMAT_DATE/FORMAT_TIMESTAMP for date formatting
+6. Use DATE_DIFF, TIMESTAMP_DIFF for date differences
+7. Use QUALIFY clause for window function filtering
+8. Use Standard SQL (prefix with #standardSQL if needed)
+"""
+
+        generated_query = await self._call_llm_for_query_generation(
+            f"{prompt}\n\n{bigquery_instructions}"
+        )
+
+        return self._ensure_bigquery_compatibility(generated_query)
+
+    def _ensure_bigquery_compatibility(self, query: str) -> str:
+        """Post-process query to ensure BigQuery compatibility."""
+        result = query
+
+        # Quote project.dataset.table names
+        result = re.sub(
+            r'(?<![`\w])(\w+)\.(\w+)\.(\w+)(?![`\w])',
+            r'`\1.\2.\3`',
+            result
+        )
+
+        # Replace NOW() with CURRENT_TIMESTAMP()
+        result = re.sub(r'\bNOW\s*\(\s*\)', 'CURRENT_TIMESTAMP()', result, flags=re.IGNORECASE)
+
+        # Replace GETDATE()
+        result = re.sub(r'\bGETDATE\s*\(\s*\)', 'CURRENT_TIMESTAMP()', result, flags=re.IGNORECASE)
+
+        return result
+
+    async def _validate_bigquery_query(self, query: str) -> QueryValidationResult:
+        """Validate BigQuery query for syntax, security, and performance."""
+        validation_result = QueryValidationResult(
+            is_valid=True,
+            query_type=None,
+            affected_tables=[],
+            estimated_cost=None,
+            warnings=[],
+            errors=[],
+            security_checks={}
+        )
+
+        try:
+            query_upper = query.strip().upper()
+
+            # Remove SQL dialect prefix
+            if query_upper.startswith('#'):
+                newline_idx = query.find('\n')
+                if newline_idx > 0:
+                    query_upper = query[newline_idx:].strip().upper()
+
+            # Determine query type
+            if query_upper.startswith(('SELECT', 'WITH')):
+                validation_result.query_type = QueryType.SELECT
+            elif query_upper.startswith('MERGE'):
+                validation_result.query_type = QueryType.UPDATE
+            else:
+                for qt in QueryType:
+                    if query_upper.startswith(qt.value):
+                        validation_result.query_type = qt
+                        break
+
+            # Extract tables
+            validation_result.affected_tables = self._extract_bigquery_tables(query)
+
+            # Security checks
+            validation_result.security_checks = {
+                "no_sql_injection_patterns": self._check_sql_injection_patterns(query),
+                "no_dangerous_operations": self._check_dangerous_bigquery_operations(query),
+                "proper_quoting": self._check_proper_quoting(query)
+            }
+
+            if not all(validation_result.security_checks.values()):
+                validation_result.is_valid = False
+                for check, passed in validation_result.security_checks.items():
+                    if not passed:
+                        validation_result.errors.append(f"Security check failed: {check}")
+
+            # Syntax validation
+            if not await self._validate_syntax_bigquery(query):
+                validation_result.is_valid = False
+                validation_result.errors.append("Invalid BigQuery SQL syntax")
+
+            # Performance warnings
+            if "SELECT *" in query_upper:
+                validation_result.warnings.append(
+                    "SELECT * scans all columns - specify needed columns for cost reduction"
+                )
+
+            if 'WHERE' not in query_upper and '_PARTITIONTIME' not in query_upper:
+                validation_result.warnings.append(
+                    "Consider adding partition filter for cost reduction"
+                )
+
+            if query.strip().startswith('#legacySQL'):
+                validation_result.warnings.append("Consider migrating to Standard SQL")
+
+            return validation_result
+
+        except Exception as e:
+            validation_result.is_valid = False
+            validation_result.errors.append(f"Validation error: {str(e)}")
+            return validation_result
+
+    def _check_dangerous_bigquery_operations(self, query: str) -> bool:
+        """Check for dangerous BigQuery operations. Returns True if SAFE."""
+        query_upper = query.upper()
+
+        dangerous_patterns = [
+            r'\bDROP\s+(TABLE|SCHEMA|VIEW|FUNCTION)\b',
+            r'\bTRUNCATE\s+TABLE\b',
+            r'\bDELETE\s+FROM\s+`[^`]+`\s*$',
+            r'\bDROP\s+ALL\s+ROW\s+ACCESS\s+POLICIES\b',
+        ]
+
+        for pattern in dangerous_patterns:
+            if re.search(pattern, query_upper, re.IGNORECASE | re.DOTALL):
+                return False
+
+        return True
+
+    def _extract_bigquery_tables(self, query: str) -> List[str]:
+        """Extract table names from BigQuery query."""
+        tables = set()
+
+        # Backtick-quoted fully-qualified names
+        tables.update(re.findall(r'`([a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+)`', query))
+        tables.update(re.findall(r'`([a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+)`', query))
+
+        # Standard table references
+        tables.update(self._extract_tables_from_query(query))
+
+        return list(tables)
+
+    async def _validate_syntax_bigquery(self, query: str) -> bool:
+        """Validate BigQuery-specific query syntax."""
+        try:
+            query_stripped = query.strip()
+
+            if query_stripped.startswith('#'):
+                newline_idx = query_stripped.find('\n')
+                if newline_idx > 0:
+                    query_stripped = query_stripped[newline_idx:].strip()
+
+            query_stripped = query_stripped.rstrip(';')
+            query_upper = query_stripped.upper()
+
+            valid_starts = [
+                'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'MERGE',
+                'CREATE', 'ALTER', 'DROP', 'TRUNCATE', 'WITH',
+                'DECLARE', 'SET', 'EXECUTE', 'BEGIN', 'IF',
+                'EXPORT', 'LOAD', 'GRANT', 'REVOKE', 'ASSERT',
+            ]
+
+            first_word = query_upper.split()[0] if query_stripped else ''
+            if first_word not in valid_starts:
+                return False
+
+            # Check balance
+            if query.count('(') != query.count(')'):
+                return False
+            if query.count('`') % 2 != 0:
+                return False
+            if query.count('[') != query.count(']'):
+                return False
+            if query.count('<') != query.count('>'):
+                return False
+
+            return True
+        except Exception:
+            return False
+
+    # =========================================================================
+    # HELPER METHODS
+    # =========================================================================
+
+    def _extract_tables_from_query(self, query: str) -> List[str]:
+        """Extract table names from SQL query."""
+        tables = set()
+
+        patterns = [
+            r'\bFROM\s+([`"\[]?[\w.-]+[`"\]]?)',
+            r'\bJOIN\s+([`"\[]?[\w.-]+[`"\]]?)',
+            r'\bINSERT\s+INTO\s+([`"\[]?[\w.-]+[`"\]]?)',
+            r'\bUPDATE\s+([`"\[]?[\w.-]+[`"\]]?)',
+            r'\bDELETE\s+FROM\s+([`"\[]?[\w.-]+[`"\]]?)',
+        ]
+
+        for pattern in patterns:
+            for match in re.findall(pattern, query, re.IGNORECASE):
+                tables.add(match.strip('`"[]'))
+
+        return list(tables)
+
+    def _find_relevant_tables(
+        self,
+        schema_metadata: SchemaMetadata,
+        natural_language_query: str
+    ) -> List[Dict[str, Any]]:
+        """Find tables relevant to the natural language query."""
+        relevant_tables = []
+        query_lower = natural_language_query.lower()
+
+        keywords = set(re.findall(r'\b\w+\b', query_lower))
+
+        stop_words = {
+            'the', 'a', 'an', 'is', 'are', 'was', 'were', 'be', 'been',
+            'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would',
+            'could', 'should', 'may', 'might', 'must', 'shall', 'can',
+            'to', 'of', 'in', 'for', 'on', 'with', 'at', 'by', 'from',
+            'and', 'or', 'but', 'if', 'as', 'show', 'get', 'find', 'list',
+            'give', 'tell', 'me', 'i', 'you', 'we', 'they', 'select', 'all',
+        }
+
+        keywords = keywords - stop_words
+
+        for table in schema_metadata.tables:
+            table_name = table.get('table_name', '').lower()
+            columns = table.get('columns', [])
+
+            score = 0
+            matched_columns = []
+
+            # Check table name
+            table_words = set(re.findall(r'\w+', table_name))
+            if table_words & keywords:
+                score += 10
+
+            for keyword in keywords:
+                if keyword in table_name:
+                    score += 5
+
+            # Check columns
+            for column in columns:
+                col_name = column.get('column_name', '').lower()
+                col_words = set(re.findall(r'\w+', col_name))
+
+                if col_words & keywords:
+                    score += 3
+                    matched_columns.append(col_name)
+
+                for keyword in keywords:
+                    if keyword in col_name and col_name not in matched_columns:
+                        score += 1
+                        matched_columns.append(col_name)
+
+            if score > 0:
+                relevant_tables.append({
+                    'table_name': table.get('table_name'),
+                    'schema': table.get('schema', schema_metadata.schema_name),
+                    'columns': columns,
+                    'matched_columns': matched_columns,
+                    'relevance_score': score,
+                    'comment': table.get('comment', '')
+                })
+
+        relevant_tables.sort(key=lambda x: x['relevance_score'], reverse=True)
+        return relevant_tables[:10]
+
+    def _extract_table_relationships(
+        self,
+        schema_metadata: SchemaMetadata,
+        relevant_tables: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Extract relationships between relevant tables."""
+        relationships = []
+        relevant_table_names = {t['table_name'] for t in relevant_tables}
+
+        # From constraints
+        for constraint in schema_metadata.constraints:
+            if constraint.get('constraint_type') == 'FOREIGN KEY':
+                source_table = constraint.get('table_name')
+                target_table = constraint.get('referenced_table')
+
+                if source_table in relevant_table_names or target_table in relevant_table_names:
+                    relationships.append({
+                        'type': 'foreign_key',
+                        'source_table': source_table,
+                        'source_column': constraint.get('column_name'),
+                        'target_table': target_table,
+                        'target_column': constraint.get('referenced_column'),
+                    })
+
+        # Infer from naming conventions
+        for table in relevant_tables:
+            for column in table.get('columns', []):
+                col_name = column.get('column_name', '')
+
+                if col_name.endswith('_id'):
+                    potential_table = col_name[:-3]
+                    for pt in [potential_table, potential_table + 's', potential_table + 'es']:
+                        if pt in relevant_table_names:
+                            relationships.append({
+                                'type': 'inferred',
+                                'source_table': table['table_name'],
+                                'source_column': col_name,
+                                'target_table': pt,
+                                'target_column': 'id',
+                            })
+                            break
+
+        return relationships
+
+    def _get_query_patterns_for_tables(
+        self,
+        relevant_tables: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Generate common query patterns for relevant tables."""
+        patterns = []
+
+        for table in relevant_tables[:3]:
+            table_name = table['table_name']
+            columns = table.get('columns', [])
+
+            if columns:
+                col_list = ', '.join([c['column_name'] for c in columns[:5]])
+                patterns.append({
+                    'description': f'Select from {table_name}',
+                    'pattern': f'SELECT {col_list} FROM {table_name} WHERE ...',
+                })
+
+            if numeric_cols := [
+                c for c in columns if c.get('data_type', '').lower() in (
+                    'integer', 'int', 'bigint', 'numeric', 'decimal', 'float'
+                )
+            ]:
+                patterns.append({
+                    'description': f'Aggregate {table_name}',
+                    'pattern': f'SELECT COUNT(*), SUM({numeric_cols[0]["column_name"]}) FROM {table_name} GROUP BY ...',
+                })
+
+        return patterns
+
+    def _get_data_type_guide(self, database_flavor: DatabaseFlavor) -> Dict[str, Any]:
+        """Get data type information for database flavor."""
+        guides = {
+            DatabaseFlavor.POSTGRESQL: {
+                'string_concat': '|| operator or CONCAT()',
+                'null_handling': 'IS NULL / IS NOT NULL, COALESCE()',
+                'boolean_type': 'BOOLEAN',
+            },
+            DatabaseFlavor.MYSQL: {
+                'string_concat': 'CONCAT() function',
+                'null_handling': 'IS NULL / IS NOT NULL, IFNULL()',
+                'boolean_type': 'TINYINT(1)',
+            },
+            DatabaseFlavor.BIGQUERY: {
+                'string_concat': 'CONCAT() or ||',
+                'null_handling': 'IS NULL / IS NOT NULL, IFNULL(), COALESCE()',
+                'boolean_type': 'BOOL',
+            }
+        }
+        return guides.get(database_flavor, guides[DatabaseFlavor.POSTGRESQL])
+
+    def _build_query_generation_prompt(
+        self,
+        natural_language: str,
+        schema_context: Dict[str, Any],
+        dialect: str
+    ) -> str:
+        """Build prompt for LLM query generation."""
+        prompt_parts = [
+            f"Generate a {dialect.upper()} SQL query for:",
+            f"\nRequest: {natural_language}",
+            f"\nDatabase: {schema_context.get('database_flavor', dialect).upper()}",
+            "\n\nAvailable Tables:",
+        ]
+
+        for table in schema_context.get('relevant_tables', [])[:5]:
+            prompt_parts.append(f"\n\nTable: {table.get('table_name')}")
+            columns = table.get('columns', [])[:15]
+            if columns:
+                prompt_parts.append("\nColumns:")
+                for col in columns:
+                    col_info = f"  - {col.get('column_name')}: {col.get('data_type', 'unknown')}"
+                    prompt_parts.append(col_info)
+
+        relationships = schema_context.get('table_relationships', [])
+        if relationships:
+            prompt_parts.append("\n\nRelationships:")
+            for rel in relationships[:5]:
+                prompt_parts.append(
+                    f"  - {rel['source_table']}.{rel['source_column']} -> "
+                    f"{rel['target_table']}.{rel['target_column']}"
+                )
+
+        prompt_parts.append("\n\nGenerate only the SQL query, no explanations:")
+        return '\n'.join(prompt_parts)
+
+    async def _call_llm_for_query_generation(self, prompt: str) -> str:
+        """Call LLM client to generate SQL query."""
+        system_msg = "You are a SQL expert. Generate precise SQL queries. Return only the SQL, no explanations."
+
+        if self.llm:
+            response = await self.llm.ask(prompt, system_prompt=system_msg)
+            if isinstance(response, AIMessage):
+                return str(response.output).strip()
+            # Handle possible dict response if client doesn't return AIMessage (fallback)
+            elif isinstance(response, dict) and 'content' in response:
+                 # Should extract text
+                 return str(response['content']).strip() # Simplified fallback
+            return str(response).strip()
+
+        if hasattr(self, 'agent') and self.agent:
+            response = await self.agent.llm.acomplete(prompt, system_message=system_msg)
+            return response.strip()
+
+        if hasattr(self, 'llm_client') and self.llm_client:
+            response = await self.llm_client.acomplete(prompt)
+            return response.strip()
+
+        raise ValueError(
+            "No LLM client configured. Provide an 'llm', 'agent' or 'llm_client' to DatabaseTool."
+        )
+
+    async def _update_schema_knowledge_base(
+        self,
+        schema_metadata: SchemaMetadata
+    ) -> None:
+        """Update knowledge store with schema metadata for RAG."""
+        if not self.knowledge_store:
+            return
+
+        documents = []
+
+        for table in schema_metadata.tables:
+            table_name = table.get('table_name')
+            columns = table.get('columns', [])
+
+            column_descriptions = [
+                f"{col.get('column_name')} ({col.get('data_type', 'unknown')})"
+                for col in columns
+            ]
+
+            doc_text = f"""
+Table: {schema_metadata.schema_name}.{table_name}
+Database: {schema_metadata.database_flavor.value}
+Columns: {', '.join(column_descriptions)}
+"""
+
+            documents.append({
+                'content': doc_text,
+                'metadata': {
+                    'type': 'database_schema',
+                    'schema': schema_metadata.schema_name,
+                    'table': table_name,
+                    'database_flavor': schema_metadata.database_flavor.value,
+                }
+            })
+
+        await self.knowledge_store.add_documents(documents)
+
+    def _convert_to_structured_output(
+        self,
+        raw_results: Any,
+        schema: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
+        """Convert raw query results to structured output."""
+        if not raw_results:
+            return []
+
+        field_mappings = schema.get('field_mappings', {})
+        structured_results = []
+
+        for row in raw_results:
+            structured_row = {}
+
+            if isinstance(row, dict):
+                if field_mappings:
+                    for target, source in field_mappings.items():
+                        if source in row:
+                            structured_row[target] = row[source]
+                        elif target in row:
+                            structured_row[target] = row[target]
+                else:
+                    structured_row = row
+            else:
+                fields = schema.get('fields', [])
+                for i, field in enumerate(fields):
+                    if i < len(row):
+                        structured_row[field] = row[i]
+
+            structured_results.append(structured_row)
+
+        return structured_results