如何用Python提取SQL脚本中的表名？解决跨行匹配难题

Alright, let's tackle this problem of pulling out all table names that follow FROM or JOIN in your SQL script—even when the table name isn't on the same line as the keyword. First, let's restate the problem clearly with your sample SQL:

Your Sample SQL #

select * from ( select col1 from test.test_a join test.test_a1 on a.col1 = a1.col1) a left join test.test_b b on a.col1 = b.col2 left join test.test_c c on b.col2 = c.col3 left join (select col4 from test.test_d) d on c.col3 = d.col4

The main challenge here is that line breaks can split the keyword (FROM/JOIN) and the table name, so naive line-by-line processing won't cut it. Let's look at two solid solutions:

Solution 1: Quick Fix with Preprocessing + Regex #

If your SQL isn't overly complex (no deeply nested subqueries or edge cases), you can normalize the SQL first to eliminate line breaks, then use a regex to grab the table names. Here's how:

Code Implementation #

import re

def extract_table_names(sql):
    # First, clean up the SQL: remove all line breaks and collapse extra spaces
    normalized_sql = re.sub(r'\s+', ' ', sql.replace('\n', ' ').strip())
    
    # Regex pattern to match table names after FROM/JOIN, skipping subqueries
    # We ignore anything that starts with '(' right after the keyword (that's a subquery)
    table_pattern = re.compile(r'(?:FROM|JOIN)\s+(?!\()([a-zA-Z0-9_.]+)', re.IGNORECASE)
    
    # Grab all matches, remove duplicates, and return as a list
    matched_tables = table_pattern.findall(normalized_sql)
    return list(set(matched_tables))

# Your SQL (works even if it's multi-line!)
sample_sql = """select * from ( select col1 from test.test_a join test.test_a1 on a.col1 = a1.col1) a left join test.test_b b on a.col1 = b.col2 left join test.test_c c on b.col2 = c.col3 left join (select col4 from test.test_d) d on c.col3 = d.col4"""

result = extract_table_names(sample_sql)
print(result)
# Output: ['test.test_a', 'test.test_a1', 'test.test_b', 'test.test_c', 'test.test_d']

How This Works #

1. Normalization: By turning all line breaks into spaces and collapsing multiple spaces into one, we make the entire SQL a single "line"—so we don't have to worry about keywords and table names being split across lines.
2. Regex: The pattern looks for FROM or JOIN, skips any whitespace, then captures the table name (including schema prefixes like test.test_a). The (?!\() part ensures we don't pick up subqueries (since those start with ( right after the keyword).

Solution 2: Robust Approach with a SQL Parsing Library #

If you need to handle more complex SQL (like deeply nested subqueries, aliases, or non-standard formatting), regex can fall short. Using a dedicated SQL parsing library like sqlparse is way more reliable because it actually understands SQL syntax.

Step 1: Install the Library #

First, grab the package via pip:

pip install sqlparse

Step 2: Code Implementation #

import sqlparse
from sqlparse.sql import IdentifierList, Identifier
from sqlparse.tokens import Token

def extract_tables(sql):
    tables = set()
    parsed_sql = sqlparse.parse(sql)[0]
    
    # Helper function to walk through the parsed SQL tokens
    def traverse_tokens(token):
        if isinstance(token, IdentifierList):
            # Handle lists of identifiers (like multiple tables in a JOIN)
            for item in token.get_identifiers():
                traverse_tokens(item)
        elif isinstance(token, Identifier):
            # Check if this is a table reference (not a column or alias)
            if token.ttype in (Token.Name, None):
                tables.add(token.get_real_name())
        elif hasattr(token, 'tokens'):
            # Look for FROM/JOIN keywords, then process the next token (the table)
            for idx, sub_token in enumerate(token.tokens):
                if sub_token.ttype == Token.Keyword and sub_token.value.upper() in ('FROM', 'JOIN'):
                    if idx + 1 < len(token.tokens):
                        traverse_tokens(token.tokens[idx + 1])
    
    traverse_tokens(parsed_sql)
    return list(tables)

# Your sample SQL
sample_sql = """select * from ( select col1 from test.test_a join test.test_a1 on a.col1 = a1.col1) a left join test.test_b b on a.col1 = b.col2 left join test.test_c c on b.col2 = c.col3 left join (select col4 from test.test_d) d on c.col3 = d.col4"""

result = extract_tables(sample_sql)
print(result)
# Output: ['test.test_a', 'test.test_a1', 'test.test_b', 'test.test_c', 'test.test_d']

How This Works #

• sqlparse parses your SQL into an abstract syntax tree (AST), which means it can tell the difference between table names, columns, aliases, and subqueries—something regex can't do consistently.
• The traverse_tokens function walks through the parsed tokens, looks for FROM/JOIN keywords, then extracts the table name that follows, even if there are line breaks or nested structures in the SQL.

Which One Should You Use? #

• Regex: Great for quick, simple cases where you know your SQL won't have weird edge cases. It's fast and doesn't require installing extra libraries.
• sqlparse: Better for production code or when you need to handle varied, complex SQL. It's more maintainable and less likely to break when the SQL format changes.

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