Using Cloudflare Hyperdrive with tokio-postgres in Workers

Using Cloudflare Hyperdrive with tokio-postgres in Workers

Cloudflare recently announced their query caching & connection pooling solution for databases in Workers, Hyperdrive.

Whilst there's plenty of examples of using Hyperdrive with libraries like pg in the documentation, someone in the Cloudflare Developers Discord asked about usability with Rust Workers.

There are examples for using tokio-postgres with workers-rs, the Rust framework for writing Cloudflare Workers, available here. Let's take that, and adapt it to use Hyperdrive.

💡
tokio-postgres does not currently support unnamed statements, which is required since connection poolers do not support prepared statements.

We will be using a fork by devsnek as the pull request has not been merged yet.

At the time of writing, it is necessary to use this fork:

[dependencies]
tokio-postgres = { git = "https://github.com/devsnek/rust-postgres/", branch = "unnamed-statement", features = ['js'], default-features = false }

The primary issue is that the Hyperdrive binding hasn't been added to the workers-rs framework, but luckily there is a generic get_binding method which we can use to cast to our Hyperdrive struct. Let's take a look at that struct:

use js_sys::Object;
use wasm_bindgen::prelude::*;
use worker::*;

#[wasm_bindgen]
extern "C" {
    #[wasm_bindgen(extends = Object)]
    pub type EdgeHyperdrive;

    #[wasm_bindgen(method, getter, js_name=connectionString)]
    pub fn connection_string(this: &EdgeHyperdrive) -> String;
}

pub struct Hyperdrive {
    inner: EdgeHyperdrive,
}

impl Hyperdrive {
    pub fn connection_string(&self) -> String {
        self.inner.connection_string()
    }
}

impl EnvBinding for Hyperdrive {
    const TYPE_NAME: &'static str = "Hyperdrive";
}


impl JsCast for Hyperdrive { ... }

impl AsRef<JsValue> for Hyperdrive { ... }

impl From<Hyperdrive> for JsValue { ... }
src/hyperdrive.rs (some impls removed for brevity)

The EnvBinding trait is primarily a TYPE_NAME, which is the name of the constructor used for sanity-checking that you're not casting a binding to the wrong type, as well as satisfying the JsCast trait.

To create a Hyperdrive binding, make sure you have at least Wrangler 3.10.1:

wrangler hyperdrive create test --connection-string="postgres://..."

With this struct and our binding ID from Wrangler, we can access our Hyperdrive binding like so:

[[hyperdrive]]
binding = "HYPERDRIVE"
id = "..."
wrangler.toml
mod hyperdrive;

let hyperdrive = env.get_binding::<hyperdrive::Hyperdrive>("HYPERDRIVE")?;
let conn_string = hyperdrive.connection_string();

Let's move onto getting this working with the TCP sockets API and tokio-postgres.

use std::str::FromStr;
use tokio_postgres::config::Config as PgConfig;

let url = Url::parse(&conn_string)?;
let hostname = url
    .host_str()
    .ok_or_else(|| RustError("unable to parse host from url".to_string()))?;

let socket = Socket::builder().connect(hostname, 5432)?;

let config = PgConfig::from_str(&conn_string)
    .map_err(|e| RustError(format!("tokio-postgres: {e:?}")))?;

We extract the hostname so we can setup a TCP socket that we'll use later, as well as create our tokio-postgres config.

use tokio_postgres::tls as PgTls;

let (client, connection) = config
    .connect_raw(socket, PgTls::NoTls)
    .await
    .map_err(|e| RustError(format!("tokio-postgres: {e:?}")))?;

wasm_bindgen_futures::spawn_local(async move {
    if let Err(error) = connection.await {
        console_log!("connection error: {:?}", error);
    }
});

With our client and connection setup, we're ready to start querying. In my case, my Hyperdrive binding is pointing to neon so I'll be querying one of their tutorial tables.

let rows = client
    .query("SELECT * FROM playing_with_neon", &[])
    .await
    .map_err(|e| RustError(format!("tokio-postgres: {e:?}")))?;

There we go! It's a bit of work to get it working currently, but that's because none of this has been upstreamed into the workers-rs library yet.

Take a look at the finished example on GitHub.