RustのSQL実行時パラメータを指定しやすくする

はじめに

tokio_postgres でクエリを実行する時、次のように記述すると思います。

let count_per_page: i64 = 100;

let rows = client.query(
    r#"
    SELECT
    FROM
      users
    WHERE
      prefecture = $1
      AND $2 <= age
    ORDER BY
      name
    LIMIT
      $3
    OFFSET
      $3 * ($4 - 1)
    "#,
    &[
        &"tokyo",
        &0_i64,
        &count_per_page,
        &2_i64,
    ]
).await?;

ここで、年齢の上限でも検索できるようにしたくなったとします。
差分は次のようになると思います。

let count_per_page: i64 = 100;

let rows = client.query(
    r#"
    SELECT
    FROM
      users
    WHERE
      prefecture = $1
      AND $2 <= age
+     AND age < $3
    ORDER BY
      name
    LIMIT
-     $3
+     $4
    OFFSET
-     $3 * ($4 - 1)
+     $4 * ($5 - 1)
    "#,
    &[
        &"tokyo",
        &0_i64,
+       &20_i64,
        &count_per_page,
        &2_i64,
    ]
).await?;

このように、パラメータを配列の途中に挿入する場合、それ以降の番号もずらす必要があります。

（末尾に追加すれば問題ないですが、検索条件をページネーションのパラメータの前後に配置したくない）

いろいろな場所を変える必要がありますし、変え忘れるとバグるのでつらいです。

Ruby の ActiveRecord みたいに、パラメータに名前をつけて実行できるとうれしいです。

  Post.find_by_sql ["SELECT body FROM comments WHERE author = :user_id OR approved_by = :user_id", { :user_id => user_id }]

マクロを作る

named_query!(
    client,
    "SELECT :hello",
    {
        hello: "hello"
    }
)

みたいに書いたら

client.query(
    "SELECT $1",
    &[
        &"hello"
    ]
)

のように展開されると嬉しいです。

なので作りました。（ほぼ ChatGPT が作ってくれました）

#[macro_export]
macro_rules! build_query_string {
    // クエリとパラメータを受け取り、プレースホルダーに置換
    ($query:expr, { $($param_name:ident : $param_value:expr),* $(,)? }) => {{
        // クエリのパラメータ名を対応する$1, $2, $3に置き換える
        let query = String::from($query);

        // パラメータリストをタプルで管理
        let mut counter = 0;  // 初期値を0に変更し、適切に参照する
        let mut query_with_placeholders = query;
        $(
            counter += 1;  // カウンタをインクリメントしながら使う
            let placeholder = format!(":{}", stringify!($param_name)); // :name, :ageなどのパラメータ
            let pg_placeholder = format!("${}", counter); // $1, $2などの置換後プレースホルダー
            query_with_placeholders = query_with_placeholders.replace(&placeholder, &pg_placeholder);
        )*

        // プレースホルダーで置換されたクエリを返す
        query_with_placeholders
    }};
}

#[macro_export]
macro_rules! build_query_array {
    // パラメータリストを受け取り、&[&dyn ToSql] 配列に変換
    ({ $($param_name:ident : $param_value:expr),* $(,)? }) => {{
        // 値を借用して配列に格納
        &[
            $(
                &$param_value
            ),*
        ]
    }};
}

#[macro_export]
macro_rules! named_query {
    // クエリとパラメータを受け取る
    ($client:expr, $query:expr, { $($param_name:ident : $param_value:expr),* $(,)? }) => {{
        // 実行するクエリを返す
        $client.query(
            &build_query_string!($query, { $($param_name : $param_value),* }),
            build_query_array!({ $($param_name : $param_value),* })
        )
    }};
}

宣言的マクロを使って、コードを生成しています。クエリの名前付きパラメータを番号に置換するのは実行時ですが、パラメータを指定する個数はさほど多くないと判断してこのままにしています。

使い心地と差分

最初のクエリはこうなります。

named_query!(
    client,
    r#"
    SELECT
    FROM
      users
    WHERE
      prefecture = :prefecture
      AND :min_age <= age
    ORDER BY
      name
    LIMIT
      :count_per_page
    OFFSET
      :count_per_page * (:page - 1)
    "#,
    {
        prefecture: "tokyo",
        min_age: 0_i64,
        count_per_page: count_per_page,
        page: 2_i64
    }
)

年齢に上限の検索を追加した時の差分は以下になります。

named_query!(
    client,
    r#"
    SELECT
    FROM
      users
    WHERE
      prefecture = :prefecture
      AND :min_age <= age
+     AND age < :upper_age
    ORDER BY
      name
    LIMIT
      :count_per_page
    OFFSET
      :count_per_page * (:page - 1)
    "#,
    {
        prefecture: "tokyo",
        min_age: 0_i64,
+       upper_age: 20_i64,
        count_per_page: count_per_page,
        page: 2_i64
    }
)

わかりやすくなりました。

（追記）sqlx 用のマクロを作る。

最近触り始めたんですけど、sqlx いいですね。query_as マクロを使うとクエリの構文エラーを見つけてくれるのが素晴らしいです。このマクロでも、クエリは $1, $2, ... と書く必要があるので、上記と同様にマクロで解決しようと考えました。なので query_as マクロをラップするマクロを作りました。

use replace_macro::named_query_sqlx;
use sqlx::postgres::PgPoolOptions;
use sqlx::query_as;
use sqlx::Error;

struct Row {
    test: Option<String>,
}

pub async fn query_by_macro_sqlx() -> Result<Option<String>, Error> {
    // DBコネクションを生成
    let db = PgPoolOptions::new()
        .max_connections(10)
        .connect("postgres://postgres:password@localhost:5432/test_db")
        .await
        .unwrap_or_else(|_| panic!("Cannot connect to the database"));

    let row = named_query_sqlx!(Row, "SELECT :hello_world AS test FROM users", {
        hello_world: "hello_world",
    })
    .fetch_one(&db)
    .await?;

    let value = row.test;

    Ok(value)
}

tokio_postgresでのマクロと同様に出るかなと思っていたのですが、ハマりました。
sqlx のマクロでは、クエリを文字列リテラルで受け取ります。なので、パラメータ名を番号に置換する処理をコンパイル時に実行する必要があります。先ほどは実行時でに置換する実装になっていたので、そのままでは使えませんでした。
文字列を置換して実行する処理を手続きマクロによって実装する必要があります。

カスタムのパーサーを作る

マクロの入力をパースする処理を記述します。

どこに責務を持たせるか迷いましたが、

• クエリで参照していないパラメータが渡されていたら無視する
• クエリで参照しているパラメータが渡されていなかったらエラーにする
  を実装しています。

lib.rs

// カスタムのパーサを定義
struct MacroInput {
    t: Type,
    original_query: LitStr,
    args: Vec<(Ident, Expr)>,
}

impl Parse for MacroInput {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let span = input.span();

        let t: Type = input.parse()?;

        input.parse::<Token![,]>()?;

        let query: LitStr = input.parse()?;

        let mut args = vec![];

        // クエリで使われているパラメータを抜き出す
        // regex が負の先読みに対応していないようなので、`::`を削除して抽出する
        let arg_name_regex = Regex::new(r":([a-zA-Z][a-zA-Z0-9_]*)").unwrap();
        // マッチした部分をベクターに収集
        let referenced_arg_names: Vec<String> = arg_name_regex
            .captures_iter(&query.value().replace("::", ""))
            .map(|cap| cap[1].to_string()) // マッチした部分全体を取得
            .collect();

        if !input.is_empty() {
            input.parse::<Token![,]>()?;

            let args_token;
            syn::braced!(args_token in input);

            let input = args_token;

            while !input.is_empty() {
                // 文字列リテラルをパース
                let arg_name: Ident = input.parse()?;

                input.parse::<Token![:]>()?;

                let arg_value: Expr = input.parse()?;

                // 次がカンマならそれを消費する（無ければ終了）
                if input.peek(Token![,]) {
                    input.parse::<Token![,]>()?;
                }

                // クエリで使われているパラメータの場合にのみ push する
                if referenced_arg_names.contains(&arg_name.to_string()) {
                    args.push((arg_name, arg_value));
                }
            }
        }

        // 指定していないパラメータがクエリで参照されていたらエラーにする
        let known_parameters = args.iter().map(|arg| arg.0.to_string()).collect_vec();
        let unknown_parameters = referenced_arg_names
            .iter()
            .filter(|arg| !known_parameters.contains(arg))
            .sorted()
            .collect_vec();

        if !unknown_parameters.is_empty() {
            return Err(syn::Error::new(
                span,
                format!(
                    "missing parameters:\n{}",
                    unknown_parameters
                        .iter()
                        .map(|name| format!("- {name}"))
                        .join("\n")
                ),
            ));
        }

        Ok(MacroInput {
            t,
            original_query: query,
            args,
        })
    }
}

マクロを記述する

lib.rs

#[proc_macro]
pub fn named_query_sqlx(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let parsed = syn::parse_macro_input!(input as MacroInput);

    let t = parsed.t;
    let arg_names = parsed
        .args
        .iter()
        .map(|arg| format!(":{}", arg.0))
        .collect_vec();
    let arg_values = parsed.args.iter().map(|arg| &arg.1).collect_vec();

    // クエリ中の名前付き引数を置き換える
    let mut query = parsed.original_query.value();

    for (index, arg_name) in arg_names.iter().enumerate() {
        query = query.replace(arg_name, &format!("${}", index + 1));
    }

    let res_token = quote! {
        query_as!(
            #t,
            #query,
            #(
                #arg_values,
            )*
        )
    };

    res_token.into()
}