1. はじめに

こんにちは、ホクソエムサポーターの藤岡です。初稿では一回で終わらせる予定だったはずの本記事もついに第三回。ついに最後です。ここまででsqlparseと構文解析の基本的な部分を解説したので、いよいよ本格的に構文解析の結果をしっかりと使うプログラムを作っていきます。

今回はsqlparseの紹介というよりは、構文規則をどうやってPythonプログラムに落とし込むか、という問題に対する自分なりの一解答例です。もっと賢いやり方はあると思いますし、もしご存知の方がいたら、ぜひコメントでご教示いただければ幸いです。

2. 注意

本記事に書かれた内容は自分の理解に基づいたものであり、誤りが含まれている可能性がありますが、ご了承ください。
もしそういった不備にお気付きの際には、コメントでご指摘いただければ幸いです。
また、以下の解説ではSQLが何度か登場しますが、すべてHiveQLです。
今回のサンプルプログラムは説明用に作成したものであり、不具合等が含まれる可能性が多分にあります。
リポジトリに入っているコードとはコメントの内容等を一部改変している部分があります。

3 サンプルプログラム: TableGraph

今回作成するのは、構文木を走査しながらテーブル/サブクエリ間の依存関係をグラフとして生成するプログラムです。

例えば、

SELECT t3.col
FROM (
  SELECT a+b AS col
  FROM t1
  UNION ALL
  SELECT c+d AS col
  FROM t2
) t3

というSQLクエリから、

クエリ
└─ t3
   ├─ t1
   └─ t2

というグラフを書きます。

実装は以下のリポジトリにあります。

github.com

ただ、今回のサンプルプログラムは行数が前回よりも少しだけ多いため、重要な箇所のみの解説とさせていただきました。代わりに、プログラムを動かして遊べるように簡単なインターフェースを実装したので、適当にprint文を差し込みながら動きを見るなどして、色々と学んでいただければ幸いです。

3.1 概要

要件は以下の通りです。

入力は1つのDMLクエリ。
入力にはCTE (With節) は含まれない。*1
出力はエッジを表すタプル（始点、終点）の集合。
エッジの始点・終点はテーブル/サブクエリ名の文字列。
クエリ全体は"__root__"という文字列で表す。
無名のサブクエリは識別できるようにIDを振る。

上の例では、3つのタプル ("__root__", "t3"), ("t3", "t1"), ("t3", "t2") からなる集合*2が得られればOKです。

グラフィカルにしたい場合はnetworkx等を使うのがいいかと思います。

3.2 実装方針

今回のプログラム作成においてポイントとなるのが、トークンとHiveQLの構文規則とをどう結びつけるか、という点です。

サンプルプログラムの主な処理はテーブル名の探索ですが、その達成には現在走査している部分がSELECT節なのかFROM節なのか、といった情報の読み取りが必要です。

こうした情報は非終端記号と呼ばれる記号で表記されます。これは以下の構文規則における、 table_reference や table_factor のことです。

table_reference:
    table_factor
  | join_table

一方、これらの非終端記号とsqlparseのトークンとは1対1で対応するものではありません。そもそもsqlparseを用いて得られる構文木は、あるSQL方言の構文規則を完全に表現したものというより、対応している各方言をだいたい全部包含したような、どっちつかずな構文木です*3。

なので、この構文木に対してさらに解析を加える必要があります。このタスクに対するアプローチは自分の思いつく限りでは以下の二つです。

構文木（もしくはその一部）をHiveQLの構文規則と対応するものに書き換える。
構文木を走査して必要な情報を探索し、集約する。

今回は、~~1のアプローチに最初気づかなかったため~~諸事情により2のアプローチを採用しました*4。基本的にHiveQLの構文規則にある各種非終端記号をクラスを使って表現し、そのクラスを用いて根トークンから走査していく方法で実装を進めます。

例えば、table_referenceをTableReferenceクラス、table_factorをTableFactorクラスによって表現していきます。

これらのクラスは、以下のラッパクラスを基底としたクラスです。

class HQLTokenWrapper:
    """
    HiveQLの構文ルールを適用するためのトークンラッパの基底クラス。
    あるトークンオブジェクトを対応する構文規則でラップしている。
    """

    def __init__(self, token: TokenType):
        if token is None:
            raise ValueError(token)
        self.token = token

    def traverse(self) -> Generator["HQLTokenWrapper", None, None]:
        """構文ルールを適用して得られるトークンをyieldするメソッド"""
        yield from []

    def nexts(self) -> List["HQLTokenWrapper"]:
        """1回のtraverseの結果得られる全てのトークンのリスト"""
        return list(self.traverse())

    @property
    def text(self) -> str:
        """トークンの文字列"""
        return str(self.token)

    def __str__(self):
        """オブジェクト情報（主にデバッグ用）"""
        clsname = self.__class__.__name__
        statement = re.sub("\n", " ", str(self.token).strip())
        if len(statement) > 10:
            return "<{} \"{}...\">".format(clsname, statement[:10])
        return "<{} \"{}\">".format(clsname, statement)

基本的には、traverseメソッドでtoken属性にあるトークンを解析し、その子孫のトークンをトークンラッパでラップして、そのtraverseをまた呼ぶ......ということを繰り返します。

といってもイメージしづらいと思うので、まずは簡単な例から順に実装を見ていきましょう。

3.3 実装解説

3.3-a. Statementトークンの中から、SELECTトークンを全て抜き出す。

クエリのルートに当たるQueryオブジェクトについて見ていきます。

QUERY_IDENTIFIER = "__root__"

class Query(HQLTokenWrapper):
    """クエリと対応するトークンのラッパ"""

    def yield_edges(self) -> Generator[Tuple, None, None]:
        """エッジを生成する"""
        token_stack = self.nexts()
        ident_stack = [(self.get_identifier(), 0)]
        while len(token_stack):
            token = token_stack.pop()
            if len(token_stack) < ident_stack[-1][1]:
                ident_stack.pop()
            if isinstance(token, (TblName, Query)):
                yield ident_stack[-1][0], token.get_identifier()
            if isinstance(token, Query):
                ident_stack.append((token.get_identifier(), len(token_stack)))
            token_stack.extend(token.nexts())

    def get_identifier(self) -> str:
        """クエリ全体に対する識別子として便宜的にQUERY_IDENTIFIERを割り当てる"""
        return QUERY_IDENTIFIER

    def traverse(self):
        """全てのSELECTトークンを抜き出す"""
        for t in self.token:
            if t.match(DML, "SELECT"):
                yield Select(t)

traverseメソッドの中身は非常にシンプルで、self.tokenの子トークンの中からSELECTトークンを取り出して、トークンラッパSelect（定義は後ほど紹介します）でラップしてyieldしているだけです。このように、構文木を部分的に走査しながら、トークンに非終端記号を当てはめていく処理です。

3.3-b. SELECTトークンの兄弟からtable_referenceとwhere_conditionを探す。

実装に入る前に、まずはSELECT節の構文規則のうち今回関係する部分について見ていきます。

SELECT [ALL | DISTINCT] select_expr, select_expr, ...
FROM table_reference
[WHERE where_condition]
[GROUP BY col_list]
[ORDER BY col_list]
[CLUSTER BY col_list
| [DISTRIBUTE BY col_list] [SORT BY col_list]
]
[LIMIT [offset,] rows]

今回の探索で重要となるテーブル名はtable_reference, where_conditionの二箇所に含まれます。これら2つを抽出するルールを書いていくのですが、ロジック自体はものすごく単純です。

SELECTトークンの兄弟を走査してFROM節の位置を特定し、そのうちのFROMトークンより後の部分を抜きだす。
WHEREトークンを抜き出す。

2については、WHERE節自体がWHEREトークンとしてまとまるように実装されているため、これだけで取り出すことができます。 1については範囲の特定が必要ですが、FROM節は始点も終点も簡単に判定できます。

以下、実装です。

class Select(HQLTokenWrapper):
    FROM_END_KEYWORD = [
        "GROUP",
        "ORDER",
        "CLUSTER",
        "DISTRIBUTE",
        "SORT",
        "LIMIT",
        "^UNION"
    ]

    @classmethod
    def is_from_end_keyword(cls, token):
        if isinstance(token, Where):
            return True
        return any(token.match(Keyword, kw, regex=True) for kw in
                   cls.FROM_END_KEYWORD)

    def traverse(self):
        """
        以下のルールに従い、table_referenceとwhere_conditionを抜き出す。
        [WITH CommonTableExpression (, CommonTableExpression)*]
        SELECT [ALL | DISTINCT] select_expr, select_expr, ...
          FROM table_reference
          [WHERE where_condition]
          [GROUP BY col_list]
          [ORDER BY col_list]
          [CLUSTER BY col_list
            | [DISTRIBUTE BY col_list] [SORT BY col_list]
          ]
         [LIMIT [offset,] rows]
        """
        token = self.token
        # UNION以降は別のSELECT節に当たるので探索範囲から外す。
        while token and not token.match(Keyword, "^UNION", regex=True):
            if token.match(Keyword, "FROM"):
                token_first_id = self.token.parent.token_index(token) + 1
                token = get_token_next(self.token.parent, token)
                # Select.FROM_END_KEYWORDでFROM節の終わりを判定する
                token = self.token.parent.token_matching(
                    self.is_from_end_keyword,
                    self.token.parent.token_index(token)
                )
                if token is None:
                    token_last = self.token.parent.tokens[-1]
                    yield TableReference.from_grouping(
                        self.token.parent,
                        token_first_id,
                        self.token.parent.token_index(token_last)
                    )
                    return
                else:
                    yield TableReference.from_grouping(
                        self.token.parent,
                        token_first_id,
                        self.token.parent.token_index(token) - 1
                    )
                    continue
            if isinstance(token, Where):
                yield WhereCondition(token)
                return
            token = get_token_next(self.token.parent, token)

traverseメソッドの他に、クラスメソッドSelect.is_from_end_keywordが定義されていますが、これはFROM節の終端を特定するためのものです。token.matchメソッドを呼び出すのが大まかな処理内容です。ただし、ここではマッチングに正規表現を使い、"^UNION"パターンでUNION, UNION ALLの両方とマッチするようにしています。

また、この方法ではマッチできないWHEREだけは別でマッチさせています。 WHEREトークン以外のトークンについては、ttype属性による判定が可能なのですが、WHEREについてはttype属性による判定ができないオブジェクト（sqlparse.sql.Whereオブジェクト）なので、matchメソッドが使えません。これはWhere以外のいくつかのトークンについても同様なのですが、どちらのケースなのかは基本的にはインポート元から判別できます。

sqlparse.sql: isinstanceによる判定
sqlparse.tokens, sqlparse.keywords: token.ttype を用いた判定もしくは token.matchを呼び出して判定

という認識で問題ないはずです。念のため、Whereのようなケースについては該当するトークンを本記事の末尾に掲載しておきます。

では、traverseメソッドについても見ていきます。まず、以下の部分はFROM節の終端に当たるトークンを探索しているコードです。

token_first_id = self.token.parent.token_index(token) + 1
token = get_token_next(self.token.parent, token)
# Select.FROM_END_KEYWORDでFROM節の終わりを判定する
token = self.token.parent.token_matching(
    self.is_from_end_keyword,
    self.token.parent.token_index(token)
)

get_token_next関数は、第一引数で渡したトークンの子の中から第二引数で渡したトークンから見て、コメントや空白をスキップした上で次のトークン（無ければNone）を返します。実装の解説は省略しますが、気になる方はこちらをどうぞ。

あとは、token.token_matchingメソッドの第一引数にis_from_end_keywordを、第二引数に探索開始地点のインデックスを渡せば目的のトークンが探索できます。

3.3-c. SELECTトークンの兄弟からtable_reference部分を抜き出す。

さて、ここまででtable_referenceは簡単に抜き出すことができましたが、ここからは少し複雑なことをしていく必要があります。

実は、3.3-bのようなシンプルな方法は構文規則に循環が存在するとうまく動作せず無限ループに入る場合があります。循環しているとは、例えば、table_referenceの構文規則をたどっていく途中のどこかでtable_referenceが左辺に出現した、というような状況です。この循環は左再帰と呼ばれます。

今回はこの左再帰が発生しているので、アプローチを変えます。自分が思いついたのは2つのアプローチです。どちらの方法も共通して、table_referenceに当たる複数のトークンを束ねる新しいトークンを定義します。

一つ目のアプローチでは、token.group_tokensメソッドを使います。

これは、複数のトークンを束ねて、grp_cls引数で指定したトークンクラスをそれらトークンの親としてインスタンス化するというメソッドです。束ねるトークンはあるトークン列の部分列でなくてはいけません。言い換えると、兄弟関係にないトークンどうしや、隣どうしでないトークンどうしを束ねることはできません。束ねる対象は、始点と終点のインデックスで指定します。

というわけで、インデックスは既に取得できる状態なので、grp_clsを用意します。複数のトークンを束ねたトークンなので、sqlparse.sql.TokenListを継承して作成します。

from sqlparse.sql import TokenList


class TableReferenceToken(TokenList):
    pass

振る舞いを追加しないので作る意味がなさそうですが、今後の拡張性や、エラートラッキングのやりやすさ、可読性を考えるとこのように定義しておいたほうがいいと思います。

話を戻すと、トークンを実際に束ねるのが以下のコードです。

token_last = self.token.parent.group_tokens(
    TableReferenceToken, 
    token_first_id, 
    self.token.parent.token_index(token) - 1
)
yield TableReference(token_last)

tokenには3.3-bで探索したトークンが入っています。実際はtokenがNullの場合も考慮しなければいけませんが、簡単なので説明は省略します。

さて、このアプローチの最大のメリットは、定義済みのメソッド（group_tokens）を使うことで実装をシンプルに済ませられることです。デメリットは、パース結果を書き換えてしまうため冪等性がなくなる可能性があったり、後続の処理に影響してしまうなどの点です。

というわけで、ここからはもう一つのアプローチ、パース結果を変更しない事例を紹介します。サンプルプログラムでもこちらの方法を採用しています。

使うのは以下の関数です。この関数はtoken.group_tokensメソッドから、元のパースツリーのトークンオブジェクトを変更する処理（+α）をごっそり削ったものです。

def group_tokens(token, grp_cls, start_idx, end_idx, include_end=True):
    """
    tokenのサブグループをgrp_clsによってまとめる。
    sqlparse純正のものから機能を大幅に少なくし、さらに元のパースツリーを書き換えないよう
    変更したもの。
    """
    end_idx = end_idx + include_end
    subtokens = token.tokens[start_idx:end_idx]
    grp = grp_cls(subtokens)
    grp.parent = token
    return grp

つまり、元のパースツリーからは繋がっていないトークンを作って、そのトークンを起点にパースツリーを葉へと掘り下げていくという方法です。

データ構造としては、グループ化するトークンとその子孫だけを切り出して有向部分木を新たに作るイメージに近いです*5。この方法の注意点は、作成したトークンの子だけは親への参照が正しくないため、 token.parent属性やそれを参照する関数等を使う場合には気をつける必要があります。

サンプルプログラムの当該箇所とは異なりますが、以下のような実装になるかと思います。

from .misc import group_tokens
token_last = group_tokens(
    token,
    TableReferenceToken, 
    token_first_id, 
    self.token.parent.token_index(token) - 1
)
yield TableReference(token_last)

3.3-d. テーブル名を取得する。

最後に、テーブル名取得についてちょっと細かい話を紹介します。なお、取得方法自体は前回の記事をご覧ください。

まず、get_aliasがWHERE IN構文に対して変な挙動を見せる点です。具体的にはWHERE col_foo IN sub_queryという構文が出現した際に、Whereオブジェクトのget_aliasメソッドを使うとcol_fooをaliasとして引っ張ってきてしまいます。

元のget_aliasメソッドの実装を読むと分かるのですが、厳密なパーサーを書いているというよりは様々な方言に広く使えるものをゆるく書いているような印象なので、バグとは言い切れないです。

なお、サンプルコードでは以下のようにsqlparseの関数をコピーしてきて少しだけ書き換えたものを実装して使っています。

class WhereCondition(HQLTokenWrapper):
    def get_subquery_alias(self):
        """
        WHERE IN 対応版のget_aliasメソッド
        """
        from sqlparse.tokens import Whitespace
        # "name AS alias"
        kw_idx, kw = self.token.token_next_by(m=(Keyword, 'AS'))
        if kw is not None:
            return self.token._get_first_name(kw_idx + 1, keywords=True)
        # "name alias" or "complicated column expression alias"
        _, ws = self.token.token_next_by(t=Whitespace)
        if len(self.token.tokens) > 2 and ws is not None:
            kw_in_idx, _ = self.token.token_next_by(m=(Keyword, "IN"))
            return self.token._get_first_name(idx=kw_in_idx, reverse=True)

次に、get_real_nameの呼び出し元についてです。テーブル名を表す最小単位のトークンはNameトークンです。しかし、Nameトークンからget_real_nameトークンを直接呼び出すとNoneが返ってきてしまいます。必ず、Identifier (Nameの場合はその親がIdentifierになっているはずです) から呼び出すようにしましょう。

サンプルプログラムの実装は以下の通りです。なお、定義されているのはTblNameというトークンラッパクラス中です。

def get_identifier(self):
    """テーブル名を取得"""
    if self.token.ttype == Name:
        return self.token.parent.get_real_name()
    if self.token.__class__ == Identifier:
        return self.token.get_real_name()

4. おわりに

これまで全3回の記事を通してsqlparseを紹介してきました。インターネット上で調べた感じだと、どうやらフォーマッタツールとして知られているようですが、その内部に定義されている種々の機能もとてもパワフルで、色々な可能性を秘めた「ライブラリ」でもあることが伝わっていれば幸いです。

ここまで色々と書いてきましたが、なんだかんだSQLは好きじゃないです。データ分析の現場では読み書きしなければいけないケースが多く、仕方なく使っているというような状態です。でも、これさえあればSQLばかりの現場でもPythonで立ち向かえるはずです。多分。

それでは、よきPythonライフを。

5. おまけ: isinstanceで判定するトークンリスト

Statement
Identifier
IdentifierList
TypedLiteral
Parenthesis
SquareBrackets
Assignment
If
For
Comparison
Comment
Where
Having
Case
Function
Begin
Operation
Values
Command

実践 Python 3

作者:Mark Summerfield
発売日: 2015/12/01
メディア: 単行本（ソフトカバー）

*1:現実に即しているとは言い難いですが、簡易化のためにこの前提を置きました

*2:サンプルコードでは可視化結果を分かりやすくするために結果の得られた順序も保持したかったので、要素が重複しないリストを返しています

*3:3章のように、そもそも間違っている例があったり、非対応のルールもあったりするのでこのように表記しました。

*4:group_tokens等のメソッドが実装されているあたりを見ると、1のアプローチの方がsqlparseの正しい使い方なのかもしれません。

*5:厳密には、これらのトークンオブジェクトの持つ親への参照を無視すれば有向部分木になります

株式会社ホクソエムのブログ

R, Python, データ分析, 機械学習

sqlparse 入門 - 応用編 -