Extract SQL Query Before Execution to Remove Unwanted Text

[Rahs0601]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Currently, when working with LLM-generated SQL queries, the responses may contain additional text such as explanations, comments, or markdown formatting (```sql ... ```). This can cause execution errors or require manual processing.

This feature proposes a method to reliably extract and clean SQL queries before execution using re and sqlparse.

Motivation

When integrating LLMs for SQL generation, the generated output often includes additional context that is not valid SQL, such as:

Markdown formatting (```sql ... ```)
Natural language explanations before or after the query

Non-SQL related text

Manually filtering these queries can be cumbersome. Automating this process improves reliability and ensures only executable SQL is returned.

Proposal (If applicable)

Implement a method like the one below to extract SQL queries reliably:

import sqlparse
import logging

logger = logging.getLogger(__name__)

def extract_query(text: str) -> str:
    """Extract SQL query from LLM response."""
    try:
        # Method 1: Extract from SQL code blocks
        sql_block_pattern = r'```(?:sql)?[\r\n]*(.*?)[\r\n]*```'
        code_block_matches = re.findall(sql_block_pattern, text, re.DOTALL | re.IGNORECASE)

        if code_block_matches:
            for match in code_block_matches:
                cleaned_query = match.strip()
                try:
                    parsed = sqlparse.parse(cleaned_query)
                    if parsed:
                        return sqlparse.format(cleaned_query, reindent=True, keyword_case='upper')
                except Exception as e:
                    logger.warning(f"Failed to parse SQL block: {str(e)}")
                    continue

        # Method 2: If no code blocks, attempt to parse the entire text as SQL
        cleaned_text = text.strip()
        return sqlparse.format(cleaned_text, reindent=True, keyword_case='upper')

    except Exception as e:
        logger.error(f"Error extracting SQL query: {str(e)}")
        return text.strip()