PG Phriday: My Postgres is Rusty

Postgres and Rust go together like peanut butter and chocolate, or ice-cream and root beer, or Batman and Robin, or mice and cheese, or sand on a beach, or crabs and elephants! Err, maybe scratch that last one.

Well regardless, there’s a whole lot of Rust going on in the Postgres world these days, especially thanks to contributions from the PGRX project. As a relative novice to Rust, I figured it was time to see what all the fuss was about and tentatively dip a foot into those turbulent and unforgiving waters.

Humble Beginnings

Check out the accompanying code to follow along.

I’ve worked on a rust project or two since working at Tembo, but there’s a lot of difference between tweaking an existing codebase and creating something from scratch. You definitely need to walk before you run with this language! It’s like C, but it’s constantly hovering over you with a ruler, ready to whack your wrist for the slightest infraction.

Thankfully there’s a lot of useful scaffolding accompanying the domineering presence. Take creating a new project, for instance:

cargo new rusty-postgres
cd rusty-postgres

Well that was easy. There are a lot of ways to interact with Postgres using Rust out there. My first instinct was to reach for the postgres crate, but apparently that’s essentially just a wrapper for tokio-postgres. Worse, getting SSL working is like pulling teeth, requiring an entirely separate postgres-openssl or postgres-native-tls crate.

Instead, I looked into what Tembo itself uses for most SQL activity: sqlx. It provides a lot of abstract wrapping to get rid of the low-level stuff, still focuses on SQL, and even provides handy tools like SQL migrations.

With a bit more experimentation for the demo, it became clear that a few other crates would also be necessary. I’d need dotenvy to easily read dev environment variables, chrono for handling dates, and rand for a bit of RNG magic. And Tokio of course, because it’s basically in everything.

Installing them is refreshingly easy, though SQLx and Tokio required a few extra features to work the way I wanted:

cargo add dotenvy chrono rand
cargo add sqlx --features "runtime-tokio tls-native-tls postgres chrono"
cargo add tokio --features "rt-multi-thread macros"

And in order to get the best mileage out of the library, it also makes sense to install the command-line component.

cargo install sqlx-cli

Once all that exists, we just need to put some code in src/main.rs and call it a day. Neat!

Bootstrapping the Schema

Before that though, it makes sense to start with some SQL. To test SQLx out a bit, let’s see how the migrate feature of the CLI tool works. We have no SQL files at all, so we need to declare one:

sqlx migrate add bootstrap-schema

And the command will respond that it created a SQL file for us to modify. Open that up in an editor and inject our schema:

CREATE TABLE sensor_log (
  id            BIGINT PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
  location      VARCHAR NOT NULL,
  reading       BIGINT NOT NULL,
  reading_date  TIMESTAMPTZ NOT NULL
);

CREATE INDEX idx_sensor_log_location ON sensor_log (location);
CREATE INDEX idx_sensor_log_reading_date ON sensor_log (reading_date);

Naturally we rely on the infamous sensor_log table used frequently in past PG Phriday blogs. before we can run this, we need one more thing: the ability to connect. So now we just need a .env file that sets the DATABASE_URL environment variable. Something like this:

DATABASE_URL=postgres://postgres:SecretPass@pg-host.com:5432

It’s important to note that Rust Postgres drivers are written from scratch and are do not link to the libpq library. As a result, some of the functionality is missing, such as the ability to use .pgpass files. So yes, that means our URL needs to include the password. Sigh.

Don’t have a Postgres instance to experiment with? Set up a Tembo Hobby Instance in a few minutes and follow along. We initially tested using Docker, but the official Postgres Docker image doesn’t support SSL, and we really wanted to test that out.

If we want to get fancy and enforce SSL this way, we need to modify the DATABASE_URL a bit to something like this:

DATABASE_URL=postgres://postgres:SecretPass@pg-host.com:5432?sslmode=verify-ca&sslrootcert=./ca.crt

Tembo has an option to download the certificate from the GUI, so that’s the easiest approach. Either URL worked in our testing.

With a connection configured, we just need to deploy the SQL itself:

sqlx migrate run

This will cause SQLx to execute any new SQL files in the database. This is actually why we love SQLx: it goes perfectly with source control. Any time we want to modify the schema, we just call sqlx migrate add and there’s a new file for our changes. The entire list of SQL migrations is applied to new environments so we can find problems early. Everything goes into source control so there are no surprises.

Getting Rusty

The Rust code itself is quite a bit more complicated than our past forays into simpler languages like Python. Though admittedly including all of the prerequisite packages is very similar:

use dotenvy::dotenv;
use rand::Rng;
use std::env;

use chrono::Local;
use std::thread::sleep;
use std::time::Duration;

use sqlx::postgres::PgPoolOptions;

Because we’re using Tokio and also relying on SQLx errors, our main function looks like this:

#[tokio::main]
async fn main() -> Result<(), sqlx::Error> {
    ...
    Ok(())
}

The code itself comes in three parts:

1. Connect to the database.
2. Insert some records.
3. View the records we’ve inserted.

Connecting is refreshingly easy. Bootstrap the environment to get our DATABASE_URL and pass it to the connect method. Of course, this being Rust, we have to also daisy-chain about 99 other calls as well. We also chose to use a connection pool because why not? It’s good practice long-term and also makes several of the subsequent function calls cleaner.

dotenv().expect(".env file not found!");
let url = env::var("DATABASE_URL").expect("DATABASE_URL not set!");

let pool = PgPoolOptions::new()
    .max_connections(5)
    .connect(&url)
    .await
    .expect("Could not connect to Postgres!");

At this point we have a choice. We can either perform our inserts directly, or start a transaction and treat it like a batch job. If we send the statements individually, each one will create an implicit transaction, so we can save a lot of overhead by creating an explicit one instead.

let mut tx = pool.begin().await?;

Then we need to loop a few inserts using simulated data. We need to initialize the random number generator, but afterwards it should output anything we need. We can use the loop counter for the sensor locations, and each row should get its own timestamp.

let mut rng = rand::thread_rng();

for i in 1..=100 {
    sqlx::query!(
        "INSERT INTO sensor_log (location, reading, reading_date)
         VALUES ($1, $2, $3)",
        &format!("A{i}"),
        rng.gen_range(1..=100),
        Local::now()
    )
    .execute(&mut *tx)
    .await?;
}

We could have just used now() directly in the INSERT, but it feels more natural for each of the “sensors” to provide its own data. Besides that, now() is pinned to the start time of the transaction, so it wouldn’t be quite accurate anyway.

The way the transaction variable gets passed is… pretty strange. If we wanted to pass the pool directly, it would be execute(&pool), so we’re just going to chalk that up to Rust doing Rust things. It probably has something to do with transactions being inherently mutable, but don’t quote us on that.

Then we just need to commit the transaction:

tx.commit().await?;

How about retrieving the data? Thankfully that’s a bit more straightforward. We can just fetch all of the rows:

let rows = sqlx::query!(
    "SELECT * 
       FROM sensor_log
      WHERE location = $1
      ORDER BY reading_date",
    "A20"
)
.fetch_all(&pool)
.await?;

And then print everything out:

for r in rows {
    println!(
        "Reading #{} at {} on {}: {}",
        r.id, r.location, r.reading_date, r.reading
    );
}

One of the handy things about SQLx is that it creates an anonymous struct based on the column names from the SQL. It’s possible to pass a predefined struct, and we probably should, but it’s possible to operate without one. More complex code that needs serde probably can’t get away with that, but that’s a task for another day.

Future Steps

We didn’t quite take every possible step in this demo. There are no updates or deletes, function calls, nothing requires triggers, and so on. This is just a bog-standard “Can we even do this at all?” kind of exploration. And as expected, the answer is “Yes”.

What could we do to enhance this or take the next step? Maybe we spawn the pool so it’s usable by an army of threads. Perhaps we should have a dedicated API and a REST interface for interacting with the table. Maybe create functions and crates for everything so we don’t just cram everything into main? The sky’s the limit!

Maybe with enough practice, we can even explore the PGRX ecosystem itself and write a very rudimentary Postgres extension. Either way, it’s certainly an interesting new world!