如何基于sqlalchemy或sqlparse判断Oracle方言SQL查询类型（含WITH语句场景）

I get it, dealing WITH clauses makes grabbing the first token useless—been there! Let's break down two solid solutions using the libraries you mentioned, starting with sqlparse since it's straightforward for token-level control.

Solution 1: Using sqlparse (Recommended for Token-Level Control) #

sqlparse lets you traverse the parsed SQL structure, so we can explicitly skip over WITH clauses and comments to find the main statement keyword. Here's a robust function:

import sqlparse
from sqlparse.tokens import Keyword, Whitespace, Comment
from sqlparse.sql import With, Statement

def determine_sql_query_type(query: str) -> str:
    # Strip leading/trailing whitespace and parse the query
    parsed = sqlparse.parse(query.strip())
    if not parsed:
        return "UNKNOWN"
    
    root_stmt = parsed[0]
    
    # Handle WITH clauses first
    if isinstance(root_stmt, With):
        # Find the main statement that follows the WITH clause
        main_stmt = root_stmt.token_next_by_instance(Statement)
        if main_stmt:
            target_stmt = main_stmt[0]
        else:
            return "UNKNOWN"
    else:
        target_stmt = root_stmt
    
    # Traverse the target statement to find the first DML/DDL keyword
    for token in target_stmt.flatten():
        if token.ttype in (Whitespace, Comment):
            continue
        if token.ttype in (Keyword.DML, Keyword.DDL):
            return token.normalized.upper()
    
    return "UNKNOWN"

How this works: #

• WITH Clause Handling: We check if the root statement is a With object, then jump directly to the main statement that comes after the CTE definitions.
• Token Filtering: We skip over whitespace and comments to avoid false positives, then grab the first DML (SELECT/UPDATE/DELETE/MERGE) or DDL (ALTER/CREATE/DROP) keyword.
• Normalization: Using normalized ensures we get consistent uppercase keywords even if the input has mixed case.

Solution 2: Using SQLAlchemy (For Dialect-Aware Parsing) #

If you prefer using SQLAlchemy's dialect-specific parser, you can leverage its internal parsing logic to get the statement type. This works well if you need Oracle-specific handling:

from sqlalchemy import text
from sqlalchemy.dialects import oracle
from sqlalchemy.sql import parsing

def get_sql_type_with_sqlalchemy(query: str) -> str:
    try:
        # Parse the SQL using Oracle dialect
        ast = parsing.parse_sql(query, oracle.dialect())
        for stmt in ast:
            # SQLAlchemy's AST nodes have a __visit_name__ attribute that maps to the statement type
            if hasattr(stmt, '__visit_name__'):
                return stmt.__visit_name__.upper()
    except Exception as e:
        # Handle invalid SQL or parsing errors
        return "UNKNOWN"
    return "UNKNOWN"

Notes on this approach: #

• Dialect Awareness: This uses Oracle's specific parser, so it handles Oracle-specific syntax better than a generic parser.
• Error Handling: We wrap in a try/except block because invalid SQL will throw a parsing error.
• AST Node Types: The __visit_name__ attribute gives us lowercase names like select, update, merge, alter, which we uppercase for consistency.

Putting It All Together with Your Existing Code #

You can integrate either function into your current workflow like this:

import sqlparse

def determine_sql_query_type(query: str) -> str:
    # Insert the function from Solution 1 here

with open(input_file_path, 'r') as file:
    raw = file.read()
    queries = sqlparse.split(raw)
    
for idx, query in enumerate(queries):
    query_type = determine_sql_query_type(query)
    print(f"Query {idx+1}: {query_type}")

This will loop through each split query and output its type without executing any SQL.

内容的提问来源于stack exchange，提问作者glezo