sqlx のマクロのために Procedural Macro を書く
ふとしです。
Procedural macro (の特に Function-like macro) を練習を兼ねて書いたので、処理の流れなどをメモとして残しておきます。
- 今回のリポジトリ: mmmpa/sqlx-concat-query
まえおき - sqlx::query_as! は便利
Rust で DB を使う時は sqlx を使っています。データを取ってくるために query_as 関数と query_as! マクロがあります。
query_as を使う場合、データをマップする構造体に FromRow を derive する必要があります。一方で query_as! はその必要がないのが大きな違いです。他にもコンパイル時にフィールドの有無のチェックが行われるなど、さまざまな強さがあります。
これは query_as! が普通の macro ではなく Procedural macro であることで実現できているのですが、一方で第 2 引数であるクエリ部は必ず裸の String literal である必要があるという縛りを生んでいます。
たとえば以下の too_huge_raw_query がどのような文字列であれクエリの再利用ができません。
sqlx::query_as!(
Record,
concat!(
too_huge_raw_query,
"where id = ?"
),
self.user_id(),
)
そこで今回はクエリ部用の String literal を好きなだけ与えられる Procedural macro を書きました。
どんな感じの macro
以下の感じで呼ぶと
concat_query_as!(
Record,
"SELECT title FROM feeds ", // ┐
"WHERE title = ? AND status = ?", // ┼─ concat
"ORDER BY title ASC". // ┘
title,
status
)
実際は以下のようになる感じです。
sqlx::query_as!(
Record,
"
SELECT title FROM feeds
WHERE title = ? AND status = ?
ORDER BY title ASC
",
title,
status
)
SELECT 部にものすごく長いクエリが入っているものがあり、WHERE だけが用法によって変わるので、再利用したいというのが今回のモチベーションです。
流れ
Procedural macro はすでに解釈された TokenStream を受けとって TokenStream を返します。返す TokenStream をどのように加工するかが各 macro の違いとなります。
受けとった TokenStream を解釈するのも、TokenStream を造るのも大変なので、それぞれに syn と quote というクレートがあります。
受けとった TokenStream を解釈する
syn の parse_macro_input! を使って、まずは扱いやすい構造体に落とし込みます。
let QueryAs {
record, // マップ先の構造体
query, // クエリ
rest, // バインドがあれば
} = parse_macro_input!(input as QueryAs);
QueryAs 構造体の定義
Derive macro を作成する場合は作り付けの DeriveInput 構造体が使えますが、Procedural macro ではおおむね自分で用意した構造体を使うことになります。今回の場合は QueryAs です。
parse_macro_input! に指定する構造体は syn::parse::Parse トレイトを満たしていればいいので、そのように定義します。
struct QueryAs {
record: Type,
query: String,
rest: TokenStream2, // proc_macro2::TokenStream
}
impl Parse for QueryAs {
fn parse(input: ParseStream) -> Result<Self> {
todo!()
}
}
あとはどのようにパースするかを書いていきます。
まずはマップ先の構造体を拾う
ParseStream から各トークンを拾うには let token:T = input.parse() (あるいは let token = input.parse::<T>()) を使います。
T になれるのは Parse を満たしているものです。(Parse#implementors)
いま現在先頭にあるトークンが T であれば Ok、そうでなければ Err を返し、ParseStream の位置が一つ進んで次のトークンを指した状態になります。
今回の macro ではマップ先の構造体が先頭に必ず入ります。従って、それ以外が入っていれば即座に終わっていいのでそのように書いていきます。
fn parse(input: ParseStream) -> Result<Self> {
let record: Type = input.parse()?;
input.parse::<Comma>()?;
todo!()
}
次は連続した String literal を拾う
前段では構造体もコンマも数が確定していましたが、クエリ部はいくつでも受け取れます。
String literal 以外のトークンが来るまですべて連結しますが、String literal 以外であることを parse を使って検査してしまうと、トークンの位置が進んでしまいます。
そこでイテレーターでもおなじみの peek 関数が用意されていますので、それを使って検査します。
fn parse(input: ParseStream) -> Result<Self> {
let record: Type = input.parse()?;
input.parse::<Comma>()?;
let mut query = String::new();
loop {
if input.peek(Comma) {
input.parse::<Comma>()?;
}
if input.peek(LitStr) {
let s: LitStr = input.parse()?;
query += &s.value();
query.push(' ');
} else {
break;
}
}
todo!()
}
コンマならば消費して、String literal なら文字列として保持、それ以外が来た場合は終わります。
残りを処理する
のこりは何でも良いので TokenStream として戻します。ただ proc_macro::TokenStream は Parse を満たしていないので、proc_macro2::TokenStream へとパースしました。
fn parse(input: ParseStream) -> Result<Self> {
let record: Type = input.parse()?;
input.parse::<Comma>()?;
let mut query = String::new();
loop {
if input.peek(Comma) {
input.parse::<Comma>()?;
}
if input.peek(LitStr) {
let s: LitStr = input.parse()?;
query += &s.value();
query.push(' ');
} else {
break;
}
}
Ok(QueryAs {
record,
query,
rest: input.parse()?,
})
}
これで macro が結果を組み立てるための要素はすべて準備できました。
解釈済みの要素を TokenStream にする。
今回は特にイテレーションする要素もないので、TokenStream を組み立てる quote! マクロに単純に渡すだけです。
#[proc_macro]
pub fn concat_query_as(input: TokenStream) -> TokenStream {
let QueryAs {
record,
query,
rest,
} = parse_macro_input!(input as QueryAs);
let tokens = quote! {
sqlx::query_as!(#record, #query, #rest)
};
tokens.into()
}
これで concat_query_as! が返した TokenStream は sqlx::query_as! を呼び出す何かとして解釈され、sqlx::query_as! が呼び出されます。
まとめ
無事生クエリ再利用が達成できました。
クエリの内容をコンパイル時にチェックしてくれる sqlx のマクロは便利なので使っていきたいですね。