Go to Qiita Advent Calendar 2024 Top
LoginSignup
12
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@nokonoko_1203(のこのこ)inMIERUNE

Rustのproc macroを書いて、構造体定義からElasticsearchのmapping定義を導出する

Last updated at Posted at 2024-03-18

弊社はGISデータ（地理空間情報）を専門的に取り扱うのですが、Elasticsearchは実はGISデータも取り扱うことができるため、検索が必要な業務システムを作成する際には頻繁に利用します。

ただ、システム構築の際に、RDBMSなどを利用する場合はORMを利用することで、プログラム内部で利用する型とデータベースのテーブル定義をマッピングさせることができますが、Elasticsearchを利用する際にはあまりそういうツールはなさそうで、型定義からElasticsearch定義に変換するためのプログラムを個別に作る必要がありました。

最近、弊社ではRustを利用していく機運が高まっているため、将来的に「構造体定義からElasticsearchのmapping定義を作成するproc macro」があると便利だろうなーと思い、ちょっと書いてみました。

今回は特定のユースケースで利用したかっただけなので、かなり雑に書いており、Rustの全ての基本型に対応しているわけではなかったり、クレートとして公開できるようなものではなかったりするのですが、備忘のために残しておきます。
（将来的にはちゃんとしたものにしていきたい。）

フォルダ構成

Rustのプロジェクトの作成方法などは他の記事に譲ります。
以下のようなフォルダ構成になりました。

❯ tree . -L 3
.
├── Cargo.lock
├── Cargo.toml
├── examples
│   └── basic_usage.rs
├── macros
│   ├── Cargo.toml
│   └── src
│       ├── lib.rs
│       └── types.rs
├── src
│   ├── lib.rs
│   └── traits
│       ├── document.rs
│       ├── mapping.rs
│       └── mod.rs
└── tests
    └── mapping.rs

14 directories, 12 files

今回はproc macroを作成していきますが、マクロの部分などはプロジェクト内部に別クレートを作成する必要があるようです。
クレート名はmacros（名称は自由）としました。

こちらの記事などを参考にしました！
[Rust] Procedural Macroの仕組みと実装方法

利用方法

以下のように、「構造体にderiveして利用する」ことを目指します。

  • examples/basic_usage.rs
use es_mapping::traits::{document::EsDocument, mapping::EsMap};
use macros::EsMapping;
use serde::{Deserialize, Serialize};
use serde_json::json;

#[derive(Serialize, Deserialize, EsMapping)]
struct User {
    #[serde(rename = "userId")]
    #[es(type = "keyword")]
    user_id: i32,
    #[es(type = "text")]
    name: String,
    #[es(type = "text")]
    age: i32,
    year: i32,
}

impl EsDocument for User {}

fn main() {
    let mapping = User::generate_mapping();
    println!("Mapping: {}", json!(mapping));
}

必要なプログラムを実装

lib.rsやmod.rsなどは省略しますので、リポジトリの方を見てください。
またCargo.tomlも同様です。

重要そうなファイルのみ記載していきます。

まずはトレイトを定義していきます。
EsMapトレイトは#[derive(EsMapping)]のように構造体にマクロが利用された場合に、その構造体に自動的に実装されるメソッドを定義します。
つまり、EsMappingをderiveするとEsMapの実装が強制され、generate_mapping()メソッドが生えます。

  • src/traits/mapping.rs
pub trait EsMap {
    fn generate_mapping() -> serde_json::Value;
}

次に、マクロ本体を定義します。

  • macros/src/lib.rs
mod types;

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{parse_macro_input, DeriveInput, LitStr};
use types::EsType;

#[proc_macro_derive(EsMapping, attributes(es))]
pub fn es_mapping_derive(token: TokenStream) -> TokenStream {
    let derive_input = &parse_macro_input!(token as DeriveInput);

    let fields = match &derive_input.data {
        syn::Data::Struct(ref data) => &data.fields,
        _ => {
            todo!("Only struct is supported");
        }
    };

    let mut es_fields: Vec<TokenStream2> = Vec::new();

    for field in fields {
        // field name
        let ident = field.ident.as_ref().unwrap();
        let ident_str = ident.to_string();

        // field type
        // todo: What to do when renaming with serde
        let ty = &field.ty;
        let es_type_text = rust_type_to_es_type(ty).as_str();
        let mut es_type = quote! { #es_type_text };

        // if the field has es attribute, override the type
        for attr in &field.attrs {
            let _ = attr.parse_nested_meta(|meta| {
                if meta.path.is_ident("type") {
                    let name: LitStr = meta.value()?.parse()?;
                    es_type = es_type_text_to_quote(name.value().as_str());
                }

                Ok(())
            });
        }
        es_fields.push(quote! {
            #ident_str: {
                "type": #es_type
            }
        });
    }

    let name = &derive_input.ident;
    let expanded: TokenStream2 = quote! {
        impl EsMap for #name {
            fn generate_mapping() -> serde_json::Value {
                serde_json::json!({
                    "mappings":{
                        "properties": {
                            #( #es_fields ),*
                        }
                    }
                })
            }
        }
    };

    TokenStream::from(expanded)
}

fn es_type_text_to_quote(es_type: &str) -> TokenStream2 {
    match es_type {
        "text" => quote! { "text" },
        "keyword" => quote! { "keyword" },
        "long" => quote! { "long" },
        "boolean" => quote! { "boolean" },
        "unsigned_long" => quote! { "unsigned_long" },
        "double" => quote! { "double" },
        "object" => quote! { "object" },
        "nested" => quote! { "nested" },
        "geo_point" => quote! { "geo_point" },
        "geo_shape" => quote! { "geo_shape" },
        _ => quote! { "text" },
    }
}

fn rust_type_to_es_type(ty: &syn::Type) -> EsType {
    match ty {
        syn::Type::Path(type_path) if type_path.path.is_ident("String") => EsType::Text,
        syn::Type::Path(type_path) if type_path.path.is_ident("bool") => EsType::Boolean,
        syn::Type::Path(type_path)
            if type_path.path.is_ident("i8")
                || type_path.path.is_ident("i16")
                || type_path.path.is_ident("i32")
                || type_path.path.is_ident("i64")
                || type_path.path.is_ident("i128")
                || type_path.path.is_ident("isize") =>
        {
            EsType::Number
        }
        syn::Type::Path(type_path)
            if type_path.path.is_ident("u8")
                || type_path.path.is_ident("u16")
                || type_path.path.is_ident("u32")
                || type_path.path.is_ident("u64")
                || type_path.path.is_ident("u128")
                || type_path.path.is_ident("usize") =>
        {
            EsType::UnsignedNumber
        }
        syn::Type::Path(type_path)
            if type_path.path.is_ident("f32") || type_path.path.is_ident("f64") =>
        {
            EsType::Double
        }
        syn::Type::Path(type_path) if type_path.path.is_ident("char") => EsType::Keyword,
        // todo: handle object
        syn::Type::Path(type_path)
            if type_path.path.is_ident("IndexMap") || type_path.path.is_ident("HashMap") =>
        {
            EsType::Object
        }
        // todo: handle nested
        syn::Type::Path(type_path) if type_path.path.is_ident("Vec") => EsType::Nested,
        syn::Type::Array(_) => EsType::Nested,
        // todo: handle Option<T>
        syn::Type::Path(type_path) if type_path.path.is_ident("Option") => EsType::Text,
        // todo: handle geo_point
        syn::Type::Path(type_path) if type_path.path.is_ident("GeoPoint") => EsType::GeoPoint,
        // todo: handle geo_shape
        syn::Type::Path(type_path) if type_path.path.is_ident("GeoShape") => EsType::GeoShape,
        _ => EsType::Text,
    }
}

基本的にマクロは「ソースコードを解析してASTに変換し、コードを付与して返す」ようなことをするため、読みづらいコードになりがちですが、今回な単純なことしかしていないので、比較的読みやすいかとは思います。
順に解説していきます。

まずes_mapping_derive関数では以下のように定義しています。

#[proc_macro_derive(EsMapping, attributes(es))]
pub fn es_mapping_derive(token: TokenStream) -> TokenStream {
    let derive_input = &parse_macro_input!(token as DeriveInput);

    let fields = match &derive_input.data {
        syn::Data::Struct(ref data) => &data.fields,
        _ => {
            todo!("Only struct is supported");
        }
    };

proc_macro_deriveマクロを利用して、マクロがどのように呼び出されるかを定義します。
今回はEsMappingという名前が定義されています。
また、attributes(es)の部分は、構造体のフィールドに対して#[es(type = "keyword")]のようにesという属性が利用できるようになることを示しています。

es_mapping_derive関数ではトークン（コンパイルされるソースコードを示すものの最小単位）のストリームを受け取り、加工されたトークンを返します。

TokenStreamがソースコード本体のようなものなので、ソースコード内に記載された構造体の情報を受け取ることができます。

ここでは、トークンから構造体かどうかチェックしています。
enumなどでも利用することができますが、今回は構造体に絞って対応しました。

    let mut es_fields: Vec<TokenStream2> = Vec::new();

    for field in fields {
        // field name
        let ident = field.ident.as_ref().unwrap();
        let ident_str = ident.to_string();

        // field type
        // todo: What to do when renaming with serde
        let ty = &field.ty;
        let es_type_text = rust_type_to_es_type(ty).as_str();
        let mut es_type = quote! { #es_type_text };

        //後述
        ...
    }

上のコードでは構造体から属性名や型などを取り出しています。
途中でrust_type_to_es_typeやes_type_text_to_quoteという関数を利用しています。

  • rust_type_to_es_type
fn rust_type_to_es_type(ty: &syn::Type) -> EsType {
    match ty {
        syn::Type::Path(type_path) if type_path.path.is_ident("String") => EsType::Text,
        syn::Type::Path(type_path) if type_path.path.is_ident("bool") => EsType::Boolean,
        syn::Type::Path(type_path)
            if type_path.path.is_ident("i8")
                || type_path.path.is_ident("i16")
                || type_path.path.is_ident("i32")
                || type_path.path.is_ident("i64")
                || type_path.path.is_ident("i128")
                || type_path.path.is_ident("isize") =>
        {
            EsType::Number
        }
        syn::Type::Path(type_path)
            if type_path.path.is_ident("u8")
                || type_path.path.is_ident("u16")
                || type_path.path.is_ident("u32")
                || type_path.path.is_ident("u64")
                || type_path.path.is_ident("u128")
                || type_path.path.is_ident("usize") =>
        {
            EsType::UnsignedNumber
        }
        syn::Type::Path(type_path)
            if type_path.path.is_ident("f32") || type_path.path.is_ident("f64") =>
        {
            EsType::Double
        }
        syn::Type::Path(type_path) if type_path.path.is_ident("char") => EsType::Keyword,
        // todo: handle object
        syn::Type::Path(type_path)
            if type_path.path.is_ident("IndexMap") || type_path.path.is_ident("HashMap") =>
        {
            EsType::Object
        }
        // todo: handle nested
        syn::Type::Path(type_path) if type_path.path.is_ident("Vec") => EsType::Nested,
        syn::Type::Array(_) => EsType::Nested,
        // todo: handle Option<T>
        syn::Type::Path(type_path) if type_path.path.is_ident("Option") => EsType::Text,
        // todo: handle geo_point
        syn::Type::Path(type_path) if type_path.path.is_ident("GeoPoint") => EsType::GeoPoint,
        // todo: handle geo_shape
        syn::Type::Path(type_path) if type_path.path.is_ident("GeoShape") => EsType::GeoShape,
        _ => EsType::Text,
    }
}

まずはrust_type_to_es_typeでRustの型を独自に定義したEsTypeという型に変換していきます。
型定義はmacros/src/types.rsで行っています。

  • macros/src/types.rs
#[derive(Debug)]
pub enum EsType {
    Text,
    Keyword,
    Number,
    Boolean,
    UnsignedNumber,
    Double,
    Object,
    Nested,
    GeoPoint,
    GeoShape,
}

impl EsType {
    pub fn as_str(&self) -> &'static str {
        match self {
            EsType::Text => "text",
            EsType::Keyword => "keyword",
            EsType::Number => "long",
            EsType::Boolean => "boolean",
            EsType::UnsignedNumber => "unsigned_long",
            EsType::Double => "double",
            EsType::Object => "object",
            EsType::Nested => "nested",
            EsType::GeoPoint => "geo_point",
            EsType::GeoShape => "geo_shape",
        }
    }
}

pub struct GeoPoint {
    pub lat: f64,
    pub lon: f64,
}

pub struct GeoShape {
    pub points: Vec<GeoPoint>,
}

EsTypeはas_str()というメソッドを持っており、Elasticsearchのマッピング定義で利用する文字列を出力します。

これを以下のようにquote!マクロでTokenStreamに変換します。
#をつけて利用することで、変数も埋め込むことができます。

        let ty = &field.ty;
        let es_type_text = rust_type_to_es_type(ty).as_str();
        let mut es_type = quote! { #es_type_text };

この処理で、Rustの型に応じたデフォルトのElasticsearchの型を定義していきました。

次のコートでは、属性ごとにElasticsearchの型を上書きするコードを書いていきます。

    let mut es_fields: Vec<TokenStream2> = Vec::new();

    for field in fields {
        // ...前述

        // if the field has es attribute, override the type
        for attr in &field.attrs {
            let _ = attr.parse_nested_meta(|meta| {
                if meta.path.is_ident("type") {
                    let name: LitStr = meta.value()?.parse()?;
                    es_type = es_type_text_to_quote(name.value().as_str());
                }

                Ok(())
            });
        }
        es_fields.push(quote! {
            #ident_str: {
                "type": #es_type
            }
        });
    }

esアトリビュートがついている場合は、es_type_text_to_quote関数を利用し、Elasticsearchの型情報を上書きしていきます。

  • es_type_text_to_quote
fn es_type_text_to_quote(es_type: &str) -> TokenStream2 {
    match es_type {
        "text" => quote! { "text" },
        "keyword" => quote! { "keyword" },
        "long" => quote! { "long" },
        "boolean" => quote! { "boolean" },
        "unsigned_long" => quote! { "unsigned_long" },
        "double" => quote! { "double" },
        "object" => quote! { "object" },
        "nested" => quote! { "nested" },
        "geo_point" => quote! { "geo_point" },
        "geo_shape" => quote! { "geo_shape" },
        _ => quote! { "text" },
    }
}

最後に、EsMapのgenerate_mappinng()メソッドを強制的に実装させるコードを書きます。

    let name = &derive_input.ident;
    let expanded: TokenStream2 = quote! {
        impl EsMap for #name {
            fn generate_mapping() -> serde_json::Value {
                serde_json::json!({
                    "mappings":{
                        "properties": {
                            #( #es_fields ),*
                        }
                    }
                })
            }
        }
    };

    TokenStream::from(expanded)

できたので、利用してみましょう。

  • examples/basic_usage.rs
use es_mapping::traits::{document::EsDocument, mapping::EsMap};
use macros::EsMapping;
use serde::{Deserialize, Serialize};
use serde_json::json;

#[derive(Serialize, Deserialize, EsMapping)]
struct User {
    #[serde(rename = "userId")]
    #[es(type = "keyword")]
    user_id: i32,
    #[es(type = "text")]
    name: String,
    #[es(type = "text")]
    age: i32,
    year: i32,
}

impl EsDocument for User {}

fn main() {
    let mapping = User::generate_mapping();
    println!("Mapping: {}", json!(mapping));
}

以下を実行すると、以下のようなマッピング定義が出てくると思います！

❯ cargo run --example basic_usage
Mapping: {"mappings":{"properties":{"age":{"type":"text"},"name":{"type":"text"},"user_id":{"type":"keyword"},"year":{"type":"long"}}}}

終わりに

ということで、Rustでマクロを書いていきました。
思ったよりも簡単にかけましたが、冒頭述べた通り、雑に書いているので、nestedやobjectなどの型にちゃんと対応しようと思うとコード量がもっと増えると思いますし、黒魔術感が出てくるのでコードリーディングも大変そうだなとは思いました。

適度に使う分には便利だと思うので、用法容量守って書いていきたいですね。

12
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
12
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?