sqlparse で SELECT 文をパースする

SELECT 文にあるエイリアスを特定したくて、sqlparse を使ってパースしてみました。

github.com

sqlparse の主要クラス

sqlparse でクエリをパースするときに、下記のクラス構造を把握しておくと挙動を掴みやすくなると思います。クラスごとの一文は、実装を読んだうえでの私の理解を記したものです。

• Statement: クエリステートメント。sqlparse.parse()の結果は list[Statement]
• Token: 自然言語処理におけるトークンのこと
• TokenList: list[Token]
• Identifier: エイリアスを付与できるトークンのまとまりを指している。TokenList の子クラス
• IdentifierList: list[Identifier]. TokenList の子クラス

こういう構造なので、パースした結果を再帰的に追っていくことで、見ようと思えば Token レベルまで細かく把握できます。

実装例

from sqlparse import parse, tokens
from sqlparse.sql import Statement, Identifier, IdentifierList, TokenList, Token

def check_token_type(token: Token, value=None, alias=None):
    identifier_list = None
    identifier = None
    identifier_value = None
    identifier_alias = None
    token_list = None

    if token.ttype in (
        tokens.Text.Whitespace,
        tokens.Text.Whitespace.Newline,
        tokens.Punctuation,
    ) or token.value in (value, alias):
        return
    elif isinstance(token, IdentifierList):
        identifier_list = token
        for identifier in identifier_list.get_identifiers():
            check_token_type(token=identifier)
    elif isinstance(token, Identifier):
        identifier = token
        identifier_value = identifier.value
        identifier_alias = identifier.get_alias()
        print(
            (
                "Identifier:: "
                f"value: {identifier_value}, "
                f"parent_name: {identifier.get_parent_name()}, "
                f"real_name: {identifier.get_real_name()}, "
                f"alias: {identifier_alias}"
            )
        )
        # for token in identifier:
        #     check_token_type(
        #         token=token, value=identifier_value, alias=identifier_alias
        #     )
    elif isinstance(token, TokenList):
        token_list = token
        for token in token_list:
            check_token_type(token=token)
    else:
        print(f"Token:: value: {token.value}, type: {token.ttype}")

if __name__ == "__main__":
    with open("test.sql") as f:
        query = f.read()

    parsed_statements: list[Statement] = parse(query)

    for parsed_statement in parsed_statements:
        for token in parsed_statement.tokens:
            check_token_type(token=token)

クエリを解析するうえでは重視しない、スペースや改行、括弧は除外しています。実行すると下記のような結果が出力されます。

Token:: value: SELECT, type: Token.Keyword.DML
Identifier:: value: TABLE_f.id, parent_name: TABLE_f, real_name: id, alias: None
Identifier:: value: TABLE_f.test_column, parent_name: TABLE_f, real_name: test_column, alias: None
Token:: value: FROM, type: Token.Keyword
Identifier:: value: PROJECT_d.DATASET_e.TABLE_f, parent_name: PROJECT_d, real_name: DATASET_e, alias: None
Token:: value: WHERE, type: Token.Keyword
Token:: value: EXISTS, type: Token.Keyword
Token:: value: SELECT, type: Token.Keyword.DML
Token:: value: *, type: Token.Wildcard
Token:: value: FROM, type: Token.Keyword
Identifier:: value: UNNEST(struct_column.array_column) AS a, parent_name: None, real_name: UNNEST, alias: a
Token:: value: WHERE, type: Token.Keyword
Identifier:: value: a.id, parent_name: a, real_name: id, alias: None
Token:: value: =, type: Token.Operator.Comparison
Token:: value: 1, type: Token.Literal.Number.Integer

エイリアスを読み取る目的では、Identifier クラスまで特定できれば十分です。Identifier.get_alias() で取得できています。

Identifier の内容をもっと細かく知りたい場合は、コメントアウトしている部分を解除すれば、再帰的に見ることができます。
これを可能にしているのは。Identifier クラスが TokenList クラスを継承しているからです。綺麗な設計だなと思いました。