"""
SQL service for SQL generation and execution logic
"""
import logging
from typing import Dict, List, Optional, Any
from repositories.postgres_repository import PostgresRepository
from services.schema_service import SchemaService
from shared.models import SQLQuery, QueryResult
from shared.exceptions import DatabaseError
logger = logging.getLogger(__name__)
class SQLService:
"""
Service for SQL generation and safe execution
Handles business logic around SQL operations
"""
def __init__(self, postgres_repo: PostgresRepository, schema_service: SchemaService, llm_config=None):
self.postgres_repo = postgres_repo
self.schema_service = schema_service
self.llm_config = llm_config
def execute_safe(self, sql: str, limit: bool = True) -> QueryResult:
"""
Execute SQL safely using the postgres repository
Args:
sql: SQL query to execute
limit: Whether to apply safety limits
Returns:
QueryResult with execution details
"""
return self.postgres_repo.execute_query(sql, limit=limit)
def get_suggested_queries(
self,
question: str,
schema_info,
limit: int = 3
) -> List[SQLQuery]:
"""
Get LLM-suggested SQL queries for a question
Args:
question: Natural language question
schema_info: Database schema information
limit: Maximum number of queries to suggest
Returns:
List of suggested SQL queries
"""
# This would be enhanced with actual LLM integration
# For now, return a placeholder
logger.info(f"Generating SQL suggestions for: {question}")
# Find relevant tables
relevant_tables = self.schema_service.find_relevant_tables(question, schema_info)
if not relevant_tables:
return []
# Generate basic queries based on question analysis
queries = []
# Count queries
if any(word in question.lower() for word in ['how many', 'count', 'total']):
for table in relevant_tables[:2]: # Limit to 2 tables
sql = f"SELECT COUNT(*) as total FROM {table}"
queries.append(SQLQuery(
sql=sql,
description=f"Count total records in {table}",
confidence=0.7
))
# List queries
if any(word in question.lower() for word in ['list', 'show', 'find', 'get']):
for table in relevant_tables[:2]:
sql = f"SELECT * FROM {table} LIMIT 10"
queries.append(SQLQuery(
sql=sql,
description=f"List records from {table}",
confidence=0.6
))
return queries[:limit]
def validate_sql(self, sql: str) -> Optional[str]:
"""
Validate SQL without executing it
Args:
sql: SQL to validate
Returns:
Error message if invalid, None if valid
"""
return self.postgres_repo.validate_tables_exist_in_sql(sql)
def explain_sql(self, sql: str) -> Dict[str, Any]:
"""
Get execution plan for SQL query
Args:
sql: SQL query to explain
Returns:
Query execution plan
"""
try:
explain_sql = f"EXPLAIN (FORMAT JSON) {sql}"
result = self.postgres_repo.execute_query(explain_sql)
if result.success and result.data:
return {
'success': True,
'plan': result.data[0].get('QUERY PLAN', [])
}
else:
return {
'success': False,
'error': result.error
}
except Exception as e:
logger.error(f"Could not explain SQL: {e}")
return {
'success': False,
'error': str(e)
}
