Tokio Postgres Dynamic Query

decapo01 included in Rust Postgresql Tokio tokio-postgres

2024-04-06 494 words 3 minutes

Contents

Note

The code can be found here.

I recently searched for this and could not find an adequate result, so I decided to make a post and a corresponding repo to demonstrate this. Sometimes you need to build a query at runtime and in doing so, you will need to create a dynamic query. You can use tokio postgres to do this but it’s not entirely straight forward.

First we want to create a generic function for handling queries in tokio postgres

use tokio_postgres::{types::ToSql, Client, Error, NoTls, Row};

async fn query<T>(
    client: &Client,
    query_str: &str,
    params: &Vec<Box<(dyn ToSql + Sync)>>,
    from_row: impl Fn(Row) -> T,
) -> Result<Vec<T>, Error> {
    let param_slice = &params.iter().map(|x| x.as_ref()).collect::<Vec<_>>();
    let rows = client.query(query_str, &param_slice).await?;
    Ok(rows.into_iter().map(from_row).collect())
}

An important note is to Box the ToSql and Sync traits. Without this, we will need lifetime parameters and it get’s a bit more complicated. The params will be the dynamic parameters we pass to the where clause of our SQL query. Note also, that this function can be used to make postgres queries in general and is not coupled with dynamic queries.

Next we’ll make an entity that maps to our table (or model) that we want to query.

#[derive(Debug)]
struct Item {
    pub id: i32,
    pub name: String,
    pub money_in_bank: i32, // just for demo, never use int for actual money
}

and we’ll need a function to map from a row to our entity

fn from_row(row: Row) -> Item {
    Item {
        id: row.get("id"),
        name: row.get("name"),
        money_in_bank: row.get("money_in_bank"),
    }
}

The next step is to create a Condition or Criteria struct to map out the possible queries that we will want to allow the system to run. Normally, I would create something called ItemConditions with this, but since I only structured this example as a cli project, I created an Args struct that’s coupled to the Clap library. I may make another post showing how to do this more generically and adding in a web framework. Nevertheless, here is what our query conditions struct looks like for this example. Note the field names.

#[derive(Parser, Debug)]
struct Args {
    #[clap(long)]
    pub id_eq: Option<i32>,

    #[clap(long)]
    pub name_eq: Option<String>,

    #[clap(long)]
    pub name_like: Option<String>,

    #[clap(long)]
    pub money_in_bank_lt: Option<i32>,

    #[clap(long)]
    pub money_in_bank_gt: Option<i32>,
}

Using the clap library means we can pass cli options to our app and it will auto serialize to this struct.

Ok great, now we need to make a function that will map these conditions to a query string and parameter vec.

fn query_param_pair_from_args(args: Args) -> (String, Vec<Box<(dyn ToSql + Sync)>>) {
    let mut query: String = "".to_string();
    let mut idx = 1;
    let mut params: Vec<Box<(dyn ToSql + Sync)>> = vec![];
    if let Some(x) = args.id_eq {
        let cond = format!("id = ${idx}");
        if query.is_empty() {
            query = cond;
        } else {
            query = format!("{query} and {cond} ");
        }
        idx = idx + 1;
        params.push(Box::new(x));
    }
    if let Some(x) = args.name_eq {
        let cond = format!("name = ${idx}");
        if query.is_empty() {
            query = cond;
        } else {
            query = format!("{query} and {cond} ");
        }
        idx = idx + 1;
        params.push(Box::new(x));
    }
    if let Some(x) = args.name_like {
        let cond = format!("name like ${idx}");
        if query.is_empty() {
            query = cond;
        } else {
            query = format!("{query} and {cond} ");
        }
        idx = idx + 1;
        params.push(Box::new(x));
    }
    if let Some(x) = args.money_in_bank_lt {
        let cond = format!("money_in_bank < ${idx}");
        if query.is_empty() {
            query = cond;
        } else {
            query = format!("{query} and {cond} ");
        }
        idx = idx + 1;
        params.push(Box::new(x));
    }
    if let Some(x) = args.money_in_bank_gt {
        let cond = format!("money_in_bank > ${idx}");
        if query.is_empty() {
            query = cond;
        } else {
            query = format!("{query} and {cond} ");
        }
        params.push(Box::new(x));
    }
    (query, params)
}

Note that as you add more field to the Args struct you’ll need to remember to add a condition in this function. You can create a sum type in Rust to represent these conditions but you would still need the struct to be serialized from the outside world. Also kinda trying to keep it simple for the example.

Now that we have all this in place we can create our main function that will run and query our db.

#[tokio::main]
async fn main() -> Result<(), Error> {
    let args = Args::parse();
    let (client, conn) =
        tokio_postgres::connect("host=localhost user=postgres password=postgres", NoTls).await?;
    tokio::spawn(async move {
        if let Err(e) = conn.await {
            eprintln!("connection error: {}", e);
        }
    });
    let (where_str, params) = query_param_pair_from_args(args);
    let base_query = "select * from items";
    let full_query = if where_str.is_empty() {
        base_query.to_string()
    } else {
        format!("{base_query} where {where_str}")
    };
    println!("{full_query}");
    let items = query(&client, &full_query.as_str(), &params, from_row).await?;
    for i in items {
        println!("{:?}", i)
    }
    Ok(())
}

That’s about it. You’ll need to create a db and populate it to run the sample and I have details for that in the repo associated with this post.