RAKUS Developers Blog | ラクス エンジニアブログ

株式会社ラクスのITエンジニアによる技術ブログです。

トップ > Go > pg_query_goではじめるSQL解析

pg_query_goではじめるSQL解析

Go PostgreSQL

はじめに

こんにちは、エンジニア３年目のTKDSです！
今回はpg_query_goについて調べてみました。
業務で使用したこともあるのですが、改めて個人的に使ってみたいと思い、使い方をさくっと調べて試しました。

pg_query_goとは

実際のPostgreSQLのSQLパーサーを使ってSQLクエリをパースして返してくれるライブラリです。
Goでも簡単にSQLの操作の取得などができてとても便利です。

github.com

簡単なサンプル

簡単なサンプルを使って、pg_query_goについて紹介します。

事前準備

事前準備をしておきます。

go mod init <プロジェクト名>

下記のサンプルコードをmain.goに保存します。

package main

import (
    "fmt"
    "log"
    "reflect"

    pg_query "github.com/pganalyze/pg_query_go/v6"
)

func main() {
    sql := `SELECT id, name FROM users WHERE active = true;INSERT INTO テーブル名 (列名1, 列名2) VALUES ('値1', '値2');`

    treeStruct, err := pg_query.Parse(sql)
    if err != nil {
        log.Fatalf("Parse failed: %v", err)
    }

    // fmt.Println(treeStruct)
    fmt.Println(treeStruct.Stmts)
    fmt.Println(len(treeStruct.Stmts)) // SQL文が増えると列数が増える
    fmt.Println(reflect.TypeOf(treeStruct.Stmts))
    fmt.Println(reflect.TypeOf(treeStruct.Stmts[0].Stmt))
    fmt.Println(reflect.TypeOf(treeStruct.Stmts[0].Stmt.Node))
    fmt.Println(reflect.TypeOf(treeStruct.Stmts[0].Stmt.Node))
    fmt.Println(treeStruct.Stmts[0].Stmt.Node)

    fmt.Println(reflect.TypeOf(treeStruct.Stmts))
    fmt.Println(reflect.TypeOf(treeStruct.Stmts[1].Stmt))
    fmt.Println(reflect.TypeOf(treeStruct.Stmts[0].Stmt.Node))
    fmt.Println(reflect.TypeOf(treeStruct.Stmts[1].Stmt.Node))
    fmt.Println(treeStruct.Stmts[1].Stmt.Node)
}

pg_query_goは、Parseメソッドで引数として受け取ったSQLを解析できます。
メソッドシグネチャはfunc pg_query.Parse(input string) (tree *pg_query.ParseResult, err error)となっており、pg_query.ParseResult型のポインタ型とエラーを返します。
pg_query.ParseResultはgithub.com/pganalyze/pg_query_go/v6@v6.1.0/pg_query.pb.goに定義されている型です。
.pb.goの拡張子なので、おそらくプロトコルバッファから生成された型のようです。

サンプルではこの構造体をたどっていくように出力しています。　　 出力の一部を記載します。
これは一番最初のfmt.Println(treeStruct.Stmts)の出力 結果です。

[stmt:{select_stmt:{target_list:{res_target:{val:{column_ref:{fields:{string:{sval:"id"}}  location:7}}  location:7}}  target_list:{res_target:{val:{column_ref:{fields:{string:{sval:"name"}}  location:11}}  location:11}}  from_clause:{range_var:{relname:"users"  inh:true  relpersistence:"p"  location:21}}  where_clause:{a_expr:{kind:AEXPR_OP  name:{string:{sval:"="}}  lexpr:{column_ref:{fields:{string:{sval:"active"}}  location:33}}  rexpr:{a_const:{boolval:{boolval:true}  location:42}}  location:40}}  limit_option:LIMIT_OPTION_DEFAULT  op:SETOP_NONE}}  stmt_len:46 stmt:{insert_stmt:{relation:{relname:"テーブル名"  inh:true  relpersistence:"p"  location:59}  cols:{res_target:{name:"列名1"  location:76}}  cols:{res_target:{name:"列名2"  location:85}}  select_stmt:{select_stmt:{values_lists:{list:{items:{a_const:{sval:{sval:"値1"}  location:102}}  items:{a_const:{sval:{sval:"値2"}  location:110}}}}  limit_option:LIMIT_OPTION_DEFAULT  op:SETOP_NONE}}  override:OVERRIDING_NOT_SET}}  stmt_location:47  stmt_len:70]

select_stmt, insert_stmtなどの記述からSELECT、INSERTなどの区別がしっかりつくように解析されているのがわかります。

さらに、

   fmt.Println(reflect.TypeOf(treeStruct.Stmts[0].Stmt.Node))
    fmt.Println(reflect.TypeOf(treeStruct.Stmts[1].Stmt.Node))

の実行結果からNodeの型がSQL操作種別を表していることがわかります。

*pg_query.Node_SelectStmt
*pg_query.Node_InsertStmt

中身をチラ見

次に中身をチラ見してみましょう。
単純に気になったのでみてみただけです。
pg_query.Parse(sql)からたどってみます。

func Parse(input string) (tree *ParseResult, err error) {
    protobufTree, err := parser.ParseToProtobuf(input)
    if err != nil {
        return
    }

    tree = &ParseResult{}
    err = proto.Unmarshal(protobufTree, tree)
    return
}

中では、parser.ParseToProtobuf、proto.Unmarshalを呼んでいるようです。
それぞれの関数の説明、ParseToProtobuf - Parses the given SQL statement into a parse tree (Protobuf format)、Unmarshal parses the wire-format message in b and places the result in m. The provided message must be mutable (e.g., a non-nil pointer to a message).
から前者はSQLをパースしてProtobuf形式で出力し、後者はそれを読めるようにして、proto.Messageに格納するようです。
まずは、parser.ParseToProtobufから見ていきます。

func ParseToProtobuf(input string) (result []byte, err error) {
    inputC := C.CString(input)
    defer C.free(unsafe.Pointer(inputC))

    resultC := C.pg_query_parse_protobuf(inputC)

    defer C.pg_query_free_protobuf_parse_result(resultC)

    if resultC.error != nil {
        err = newPgQueryError(resultC.error)
        return
    }

    result = []byte(C.GoStringN(resultC.parse_tree.data, C.int(resultC.parse_tree.len)))

    return
}

中身はCGOを呼び出すコードになってました。
CがCGOに紐付いたライブラリを呼び出しているコードのようです。
Cで扱えるようにGoの文字列を変換したあと、C.pg_query_parse_protobufで解析し、またGoで扱えるようにしていました。
ここでProtobufの形式でやりとりしていたので、後にUnmashalが必要になるのですね。
戻って、次はproto.Unmarshalです。これは受け取ったメッセージをtreeに書き込んでいます。
treeの引数の型部分にあったproto.MessageはただのInterfaceなので、実装はParseResultが持っていそうです。
ParseResultのあるファイルの中身を探しているとStringメソッドがありました。

func (x *ScanResult) String() string {
    return protoimpl.X.MessageStringOf(x)
}

今回出力されてる結果はこのメソッドの出力だと検討がつきました。
満足したのでこの辺にしておきます。
次からは簡単な実用例を示します。

実用例：SQLの操作を抽出

以下のようにパースしたあとに型の判別を追加します。
前述のようにParseResult.Stmts.RawStmt.Stmt.Nodeの型が操作種別を表しているようです。
そこで、型によって処理を分岐させることで、操作を抽出してみましょう。

package main

import (
    "fmt"
    "log"
    "reflect"

    pg_query "github.com/pganalyze/pg_query_go/v6"
)

func main() {
    sql := `SELECT id, name FROM users WHERE active = true;INSERT INTO テーブル名 (列名1, 列名2) VALUES ('値1', '値2');`
    treeStruct, err := pg_query.Parse(sql)
    if err != nil {
        log.Fatalf("Parse failed: %v", err)
    }

    var opeType string
    for _, stmt := range treeStruct.Stmts {
        node := stmt.Stmt.Node

        switch node.(type) {
        case *pg_query.Node_SelectStmt:
            opeType = "SELECT"
        case *pg_query.Node_InsertStmt:
            opeType = "INSERT"
        case *pg_query.Node_UpdateStmt:
            opeType = "UPDATE"
        case *pg_query.Node_DeleteStmt:
            opeType = "DELETE"
        default:
            opeType = "UNKOWN"
        }

        fmt.Printf("操作種別：%s, 型：%v\n", opeType, reflect.TypeOf(node))
    }
}

実行結果は以下の通りです。

操作種別：SELECT, 型：*pg_query.Node_SelectStmt
操作種別：INSERT, 型：*pg_query.Node_InsertStmt

みごと２つの文とも操作種別の抽出ができました！

実用例：SQL実行種別のガードレール

前項で操作種別の抽出が可能になりました。
ということは、許可したクエリのみ実行できるように制限することも可能です。
さらにロジックを追加します。
今回はSELECTのみ許可するようにしましょう。
fmt.Printf("操作種別：%s, 型：%v\n", opeType, reflect.TypeOf(node))
を削除して、以下のロジックを追加してください。

↓追加ロジック

if opeType == "SELECT" {
    fmt.Printf("操作種別：%s, 型：%v\n", opeType, reflect.TypeOf(node))
}

出力は以下の通りで、SELECTのときのみ出力されるようになりました。

操作種別：SELECT, 型：*pg_query.Node_SelectStmt

今回はPrint文でしたが、SQLなどの実行処理の前にこの操作チェックをおけば、実行のガードレールとして使用可能です。

まとめ

今回はpg_query_goについて紹介しました！
SQLの解析が簡単にできて便利です。
実用例の活用方法として、ユーザー入力のSQLを実行するMCPサーバーなど作る場合に使えそうです。
他にも地味にいいポイントが、Goなのでビルドしてしまえば簡単に配布ができます。
実行もgo mod tidy→go run main.goだけなので非常に試しやすいです。
ぜひ興味もった方はご自身の環境でためしてみてください。
今後は、pg_queryにWASM版もあるらしいのでそちらも試してみたいです。
ここまで読んでいただきありがとうございました！

ラクスの技術 / デザイン情報サイト
Copyright © RAKUS Co., Ltd. All rights reserved.