Photo - Jonny Gios
Photo

Forging SQL from Rust

How pgx creates PostgreSQL extension SQL definitions that bind to Rust.

Posted on 2021-09-20, around 21 minutes of reading.

PostgreSQL offers an extension interface, and it's my belief that Rust is a fantastic language to write extensions for it. Eric Ridge thought so too, and started pgx awhile back. I've been working with him to improve the toolkit, and wanted to share about one of our latest hacks: improving the generation of extension SQL code to interface with Rust.

This post is more on the advanced side, as it assumes knowledge of both Rust and PostgreSQL. We'll approach topics like foreign functions, dynamic linking, procedural macros, and linkers.

Understanding the problem

pgx based PostgreSQL extensions ship as the following:

$ tree plrust
plrust
├── plrust-1.0.sql
├── libplrust.so
└── plrust.control

These extension objects include a *.control file which the user defines, a *.so cdylib that contains the compiled code, and *.sql which PostgreSQL loads the extension SQL entities from. We'll talk about generating this SQL today!

These sql files must be generated from data within the Rust code. The SQL generator pgx provides needs to:

An earlier version of cargo-pgx had cargo-pgx pgx schema command that would read your Rust files and generate SQL corresponding to them.

This worked okay! But gosh, it's not fun to do that, and there are a lot of complications such as trying to resolve types!

So, what's more fun than parsing Rust source code and generating SQL? Parsing Rust code in procedural macros to inject metadata foreign functions, then later creating a binary which re-exports those functions via linker tricks, dynamically loads itself, and calls them all to collect metadata, then builds a depdency graph of them to drive the output!

Wait... What, that was not your answer? Oh no... Well, bear with me because that's what we're doing.

So, why else should we do this other than fun?

For more than fun

Okay, did I convince you? ... No? Dangit. Oh... well, Let's explore anyways!

Here's some fun ideas we pondered (and, in some cases, tried):

But wait! It turns out, that can work! We can have the binary re-export the functions and be very careful with what we use!

An illuminated path in a forest. - Johannes Plenio
An illuminated path in a forest.

The (longish) short of it

We're going to produce a binary during the build. We'll have macros output some foreign functions, then the binary will re-export and call them to build up a structure representing our extension. cargo-pgx pgx schema's job will be to orcestrate that.

Roughly, we're slipping into the build process this way:

                            ┌────────────────────────┐
                            │cargo pgx discovers     │
                            │__pgx_internal functions│
                            └───────────────────┬────┘
                                                │
                                                ▼
   Build.rs ─► Macros ─► Compile ─► Link ─► Discovery ─► Generation
                 ▲                   ▲                       ▲
                 │                   │                       │
┌────────────────┴─────┐  ┌──────────┴────────┐ ┌────────────┴────────────┐
│Parse definitions.    │  │Re-export internal │ │Binary recieves list,    │
│                      │  │functions to binary│ │dynamically loads itself,│
│Create __pgx_internal │  │via dynamic-list   │ │creates dependency graph,│
│metadata functions    │  │                   │ │outputs SQL              │
└──────────────────────┘  └───────────────────┘ └─────────────────────────┘

During the proc macro expansion process, we can append the proc_macro2::TokenStream with some metadata functions. For example:

#[derive(pgx::PostgresType)]
struct Floof { boof: usize }

// We should extend the TokenStream with something like
#[no_mangle]
pub extern "C" fn __pgx_internals_type_Floof()
-> pgx::datum::inventory::SqlGraphEntity {
    todo!()
}

How about functions? Same idea:

#[pg_extern]
pub fn floof_from_boof(boof: usize) -> Floof {
    Floof { boof }
}

// We should extend the TokenStream with something like
#[no_mangle]
pub extern "C" fn __pgx_internals_fn_floof_from_boof()
-> pgx::datum::inventory::SqlGraphEntity {
    todo!()
}

Then, in our sql-generator binary, we need to re-export them! We can do this by setting linker in .cargo/config to a custom script which includes dynamic-list:

# ...
[target.x86_64-unknown-linux-gnu]
linker = "./.cargo/pgx-linker-script.sh"
# ...

Note: We can't use rustflags here since it can't handle environment variables or relative paths.

In .cargo/pgx-linker-script.sh:

#! /usr/bin/env bash
# Auto-generated by pgx. You may edit this, or delete it to have a new one created.

if [[ $CARGO_BIN_NAME == "sql-generator" ]]; then
    UNAME=$(uname)
    if [[ $UNAME == "Darwin" ]]; then
	TEMP=$(mktemp pgx-XXX)
        echo "*_pgx_internals_*" > ${TEMP}
        gcc -exported_symbols_list ${TEMP} $@
        rm -rf ${TEMP}
    else
        TEMP=$(mktemp pgx-XXX)
        echo "{ __pgx_internals_*; };" > ${TEMP}
        gcc -Wl,-dynamic-list=${TEMP} $@
        rm -rf ${TEMP}
    fi
else
    gcc -Wl,-undefined,dynamic_lookup $@
fi

We also need to ensure that the Cargo.toml has a few relevant settings:

[lib]
crate-type = ["cdylib", "rlib"]

[profile.dev]
# avoid https://github.com/rust-lang/rust/issues/50007
lto = "thin"

[profile.release]
# avoid https://github.com/rust-lang/rust/issues/50007
lto = "fat"

Since these functions all have a particular naming scheme, we can scan for them in cargo pgx schema like so:

let dsym_path = sql_gen_path.resolve_dsym();
let buffer = ByteView::open(dsym_path.as_deref().unwrap_or(&sql_gen_path))?;
let archive = Archive::parse(&buffer).unwrap();

let mut fns_to_call = Vec::new();
for object in archive.objects() {
    match object {
        Ok(object) => match object.symbols() {
            SymbolIterator::Elf(iter) => {
                for symbol in iter {
                    if let Some(name) = symbol.name {
                        if name.starts_with("__pgx_internals") {
                            fns_to_call.push(name);
                        }
                    }
                }
            }
            _ => todo!(),
        },
        Err(_e) => {
            todo!();
        }
    }
}

This list gets passed into the binary, which dynamically loads the code using something like:

let mut entities = Vec::default();
// We *must* use this or the extension might not link in.
let control_file = __pgx_marker()?;
entities.push(SqlGraphEntity::ExtensionRoot(control_file));
unsafe {
    let lib = libloading::os::unix::Library::this();
    for symbol_to_call in symbols_to_call {
        let symbol: libloading::os::unix::Symbol<
            unsafe extern fn() -> SqlGraphEntity
        > = lib.get(symbol_to_call.as_bytes()).unwrap();
        let entity = symbol();
        entities.push(entity);
    }
};

Then the outputs of that get passed to our PgxSql structure, something like:

let pgx_sql = PgxSql::build(
    pgx::DEFAULT_TYPEID_SQL_MAPPING.clone().into_iter(),         
    pgx::DEFAULT_SOURCE_ONLY_SQL_MAPPING.clone().into_iter(),
    entities.into_iter()
).unwrap();

tracing::info!(path = %path, "Writing SQL");
pgx_sql.to_file(path)?;
if let Some(dot_path) = dot {
    tracing::info!(dot = %dot_path, "Writing Graphviz DOT");
    pgx_sql.to_dot(dot_path)?;
}

Since SQL is very order dependent, and Rust is largely not, our SQL generator must build up a dependency graph, we use petgraph for this. Once a graph is built, we can topological sort from the root (the control file) and get out an ordered set of SQL entities.

Then, all we have to do is turn them into SQL and write them out!

Old cogs. - Isis França
Old cogs.

Fitting the pieces together

Let's talk more about some of the moving parts. There are a few interacting concepts, but everything starts in a few proc macros.

Understanding syn/quote

A procedural macro in Rust can slurp up a TokenStream and output another one. My favorite way to parse a token stream is with syn, to output? Well that's quote!

syn::Parse and quote::ToTokens do most of the work here.

syn contains a large number of predefined structures, such as syn::DeriveInput which can be used too. Often, your structures will be a combination of several of those predefined structures.

You can call parse, parse_str, or parse_quote to create these types:

let val: syn::LitBool = syn::parse_str(r#"true"#)?;
let val: syn::LitBool = syn::parse_quote!(true);
let val: syn::ItemStruct = syn::parse_quote!(struct Floof;);

parse_quote works along with quote::ToTokens. Its use is similar to that of how the quote macro works!

use quote::ToTokens;
let tokens = quote::quote! { struct Floof; };
println!("{}", tokens.to_token_stream());
// Prints:
// struct Floof ;

They get called by proc macro declarations:

#[proc_macro_derive(Example, attributes(demo))]
pub fn example(input: TokenStream) -> TokenStream {
    let parsed = parse_macro_input!(input as ExampleParsed);
    let retval = handle(parsed);
    retval.to_token_stream()
}

In the case of custom derives, often something like this works:

#[proc_macro_derive(Example, attributes(demo))]
pub fn example(input: TokenStream) -> TokenStream {
    let parsed = parse_macro_input!(input as syn::DeriveInput);
    let retval = handle(parsed);
    retval.to_token_stream()
}

Rust's TypeIds & Type Names

Rust's standard library is full of treasures, including core::any::TypeId and core::any::type_name.

Created via TypeId::of<T>(), TypeIds are unique TypeId for any given T.

use core::any::TypeId;
assert!(TypeId::of<&str>() == TypeId::of<&str>());
assert!(TypeId::of<&str>() != TypeId::of<usize>());
assert!(TypeId::of<&str>() != TypeId::of<Option<&str>>());

We can use this to determine two types are indeed the same, even if we don't have an instance of the type itself.

During the macro expansion phase, we can write out TypeId::of<#ty>() for each used type pgx interacts with (including 'known' types and user types.)

Later in the build phase these calls exist as TypeId::of<MyType>(), then during the binary phase, these TypeIds get evaluated and registered into a mapping, so they can be queried.

core::any::type_name<T>() is a diagnostic function available in core that makes a 'best-effort' attempt to describe the type.

assert_eq!(
    core::any::type_name::<Option<String>>(),
    "core::option::Option<alloc::string::String>",
);

Unlike TypeIds, which result in the same ID at in any part of the code, type_name cannot promise this, from the docs:

The returned string must not be considered to be a unique identifier of a type as multiple types may map to the same type name. Similarly, there is no guarantee that all parts of a type will appear in the returned string: for example, lifetime specifiers are currently not included. In addition, the output may change between versions of the compiler.

So, type_name is only somewhat useful, but it's our best tool for inspecting the names of the types we're working with. We can't depend on it, but we can use it to infer things, leave human-friendly documentation, or provide our own diagnostics.

#[derive(pgx::PostgresType)]
struct Floof { boof: usize }

The above gets parsed and new tokens are quasi-quoted atop this template like this:

#[no_mangle]
pub extern "C" fn  #inventory_fn_name() -> pgx::datum::inventory::SqlGraphEntity {
    let mut mappings = Default::default();
    <#name #ty_generics as pgx::datum::WithTypeIds>::register_with_refs(
        &mut mappings,
        stringify!(#name).to_string()
    );
    pgx::datum::WithSizedTypeIds::<#name #ty_generics>::register_sized_with_refs(
        &mut mappings,
        stringify!(#name).to_string()
    );
    // ...
    let submission = pgx::inventory::InventoryPostgresType {
        name: stringify!(#name),
        full_path: core::any::type_name::<#name #ty_generics>(),
        mappings,
        // ...
    };
    pgx::datum::inventory::SqlGraphEntity::Type(submission)
}

The proc macro passes will make it into:

#[no_mangle]
pub extern "C" fn  __pgx_internals_type_Floof() -> pgx::datum::inventory::SqlGraphEntity {
    let mut mappings = Default::default();
    <Floof as pgx::datum::WithTypeIds>::register_with_refs(
        &mut mappings,
        stringify!(Floof).to_string()
    );
    pgx::datum::WithSizedTypeIds::<Floof>::register_sized_with_refs(
        &mut mappings,
        stringify!(Floof).to_string()
    );
    // ...
    let submission = pgx::inventory::InventoryPostgresType {
        name: stringify!(Floof),
        full_path: core::any::type_name::<Floof>(),
        mappings,
        // ...
    };
    pgx::datum::inventory::SqlGraphEntity::Type(submission)
}

Next, let's talk about how the TypeId mapping is constructed!

Mapping Rust types to SQL

We can build a TypeId mapping of every type pgx itself has builtin support. For example, we could do:

let mut mapping = HashMap::new();
mapping.insert(TypeId::of<T>(), RustSqlMapping {
    rust: type_name<T>(),
    sql: "MyType",
    id: TypeId::of<T>(),
});

This works fine, until we get to extension defined types. They're a bit different!

Since #[pg_extern] decorated functions can use not only some custom type T, but also other types like PgBox<T>, Option<T>, Vec<T>, or pgx::datum::Array<T> we want to also create mappings for those. So we also need TypeIds for those... if they exist.

Types such as Vec<T> require a type to be Sized, pgx::datum::Array<T> requires a type to implement IntoDatum. These are complications, since we can't always do something like that unless MyType implements those. Unfortunately, Rust doesn't really give us the power in macros to do something like what the impls crate does in macros, so we can't do something like:

// This doesn't work
syn::parse_quote! {
    #[no_mangle]
    pub extern "C" fn  #inventory_fn_name() -> pgx::datum::inventory::SqlGraphEntity {
        let mut mappings = Default::default();
        #hypothetical_if_some_type_impls_sized_block {
            mapping.insert(TypeId::of<Vec<#name>>(), RustSqlMapping {
                rust: core::any::type_name::<#name>(),
                sql: core::any::type_name::<#name>.to_string() + "[]",
                id: TypeId::of<Vec<MyType>>(),
            });
        }
        // ...   
    }
};

Thankfully we can use the same strategy as impls:

Inherent implementations are a higher priority than trait implementations.

First, we'll create a trait, and define a blanket implementation:

pub trait WithTypeIds {
    /// The [`core::any::TypeId`] of some `T`
    const ITEM_ID: Lazy<TypeId>;
    /// The [`core::any::TypeId`] of some `Vec<T>`, if `T` is sized.
    const VEC_ID: Lazy<Option<TypeId>>;
    // ...
}

impl<T: 'static + ?Sized> WithTypeIds for T {
    const ITEM_ID: Lazy<TypeId> = Lazy::new(|| TypeId::of::<T>());
    const VEC_ID: Lazy<Option<TypeId>> = Lazy::new(|| None);
    // ...
}

This lets us do <T as WithTypeIds>::ITEM_ID for any T, but the VEC_ID won't ever be populated. Next, we'll create a 'wrapper' holding only a core::marker::PhantomData<T>

pub struct WithSizedTypeIds<T>(pub core::marker::PhantomData<T>);

impl<T: 'static> WithSizedTypeIds<T> {
    pub const VEC_ID: Lazy<Option<TypeId>> = Lazy::new(||
        Some(TypeId::of::<Vec<T>>())
    );
    // ...
}

Now we can do WithSizedTypeIds::<T>::VEC_ID for any T to get the TypeId for Vec<T>, and only get Some(item) if that type is indeed sized.

Using this strategy, we can have our __pgx_internals functions build up a mapping of TypeIds and what SQL they map to.

Our pet graph

Once we have a set of SQL entities and a mapping of how different Rust types can be represented in SQL we need to figure out how to order it all.

While this is perfectly fine in Rust:

struct Dog { floof: Floofiness }

enum Floofiness { Normal, Extra }

The same in SQL is not valid.

We use a petgraph::stable_graph::StableGraph, inserting all of the SQL entities, then looping back and connecting them all together.

If an extension has something like this:

#[derive(PostgresType, Serialize, Deserialize, Debug, Eq, PartialEq)]
pub struct Animals {
    // ...
}

#[pg_extern]
fn known_animals() -> Animals {
    todo!()
}

We need to go and build an edge reflecting that the function known_animals requires the type Animals to exist. It also needs edges reflecting that these entities depend on the extension root.

Building up the graph is a two step process, first we populate it with all the SqlGraphEntity we found.

This process involves adding the entity, as well doing things like ensuring the return value has a node in the graph even if it's not defined by the user. Something like &str, a builtin value pgx knows how to make into SQL and back.

fn initialize_externs(
    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
    root: NodeIndex,
    bootstrap: Option<NodeIndex>,
    finalize: Option<NodeIndex>,
    externs: Vec<PgExternEntity>,
    mapped_types: &HashMap<PostgresTypeEntity, NodeIndex>,
    mapped_enums: &HashMap<PostgresEnumEntity, NodeIndex>,
) -> eyre::Result<(
    HashMap<PgExternEntity, NodeIndex>,
    HashMap<String, NodeIndex>,
)> {
    let mut mapped_externs = HashMap::default();
    let mut mapped_builtin_types = HashMap::default();
    for item in externs {
        let entity: SqlGraphEntity = item.clone().into();
        let index = graph.add_node(entity.clone());
        mapped_externs.insert(item.clone(), index);
        build_base_edges(graph, index, root, bootstrap, finalize);

        // ...
        match &item.fn_return {
            PgExternReturnEntity::None | PgExternReturnEntity::Trigger => (),
            PgExternReturnEntity::Iterated(iterated_returns) => {
                for iterated_return in iterated_returns {
                    let mut found = false;
                    for (ty_item, &_ty_index) in mapped_types {
                        if ty_item.id_matches(&iterated_return.0) {
                            found = true;
                            break;
                        }
                    }
                    for (ty_item, &_ty_index) in mapped_enums {
                        if ty_item.id_matches(&iterated_return.0) {
                            found = true;
                            break;
                        }
                    }
                    if !found {
                        mapped_builtin_types
                            .entry(iterated_return.1.to_string())
                            .or_insert_with(|| {
                                graph.add_node(SqlGraphEntity::BuiltinType(
                                    iterated_return.1.to_string(),
                                ))
                            });
                    }
                }
            }
            // ...
        }
    }
    Ok((mapped_externs, mapped_builtin_types))
}

Once the graph is fully populated we can circle back and connect rest of the things together! This process includes doing things like connecting the arguments and returns of our #[pg_extern] marked functions.

Something like this:

fn connect_externs(
    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
    externs: &HashMap<PgExternEntity, NodeIndex>,
    schemas: &HashMap<SchemaEntity, NodeIndex>,
    types: &HashMap<PostgresTypeEntity, NodeIndex>,
    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
    builtin_types: &HashMap<String, NodeIndex>,
    extension_sqls: &HashMap<ExtensionSqlEntity, NodeIndex>,
) -> eyre::Result<()> {
    for (item, &index) in externs {
        for (schema_item, &schema_index) in schemas {
            if item.module_path == schema_item.module_path {
                tracing::debug!(
                    from = %item.rust_identifier(),
                    to = %schema_item.rust_identifier(),
                    "Adding Extern after Schema edge"
                );
                graph.add_edge(schema_index, index, SqlGraphRelationship::RequiredBy);
                break;
            }
        }
        
        // ...
        match &item.fn_return {
            PgExternReturnEntity::None | PgExternReturnEntity::Trigger => (),
            PgExternReturnEntity::Iterated(iterated_returns) => {
                for iterated_return in iterated_returns {
                    let mut found = false;
                    for (ty_item, &ty_index) in types {
                        if ty_item.id_matches(&iterated_return.0) {
                            tracing::debug!(
                                from = %item.rust_identifier(),
                                to = %ty_item.rust_identifier(),
                                "Adding Extern after Type (due to return) edge"
                            );
                            graph.add_edge(
                                ty_index,
                                index,
                                SqlGraphRelationship::RequiredByReturn
                            );
                            found = true;
                            break;
                        }
                    }
                    if !found {
                        for (ty_item, &ty_index) in enums {
                            if ty_item.id_matches(&iterated_return.0) {
                                tracing::debug!(
                                    from = %item.rust_identifier(),
                                    to = %ty_item.rust_identifier(),
                                    "Adding Extern after Enum (due to return) edge."
                                );
                                graph.add_edge(
                                    ty_index,
                                    index,
                                    SqlGraphRelationship::RequiredByReturn,
                                );
                                found = true;
                                break;
                            }
                        }
                    }
                    if !found {
                        let builtin_index = builtin_types
                            .get(&iterated_return.1.to_string())
                            .expect(&format!(
                                "Could not fetch Builtin Type {}.",
                                iterated_return.1
                            ));
                        tracing::debug!(
                            from = %item.rust_identifier(),
                            to = iterated_return.1,
                            "Adding Extern after BuiltIn Type (due to return) edge"
                        );
                        graph.add_edge(
                            *builtin_index,
                            index,
                            SqlGraphRelationship::RequiredByReturn,
                        );
                    }
                    // ...
                }
            }
        }
    }
    Ok(())
}

With a graph built, we can topologically sort the graph and transform them to SQL representations:

#[tracing::instrument(level = "error", skip(self))]
pub fn to_sql(&self) -> eyre::Result<String> {
    let mut full_sql = String::new();
    for step_id in petgraph::algo::toposort(&self.graph, None)
        .map_err(|e| eyre_err!("Failed to toposort SQL entities: {:?}", e))?
    {
        let step = &self.graph[step_id];

        let sql = step.to_sql(self)?;

        if !sql.is_empty() {
            full_sql.push_str(&sql);
            full_sql.push('\n');
        }
    }
    Ok(full_sql)
}

Generating the SQL

In order to generate SQL for an entity, define a ToSql trait which our SqlGraphEntity enum implements like so:

pub trait ToSql {
    fn to_sql(&self, context: &PgxSql) -> eyre::Result<String>;
}

/// An entity corresponding to some SQL required by the extension.
#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub enum SqlGraphEntity {
    Schema(InventorySchema),
    Enum(InventoryPostgresEnum),
    // ...
}

impl ToSql for SqlGraphEntity {
    #[tracing::instrument(
        level = "debug",
        skip(self, context),
        fields(identifier = %self.rust_identifier())
    )]
    fn to_sql(&self, context: &super::PgxSql) -> eyre::Result<String> {
        match self {
            SqlGraphEntity::Schema(item) => {
                if item.name != "public" && item.name != "pg_catalog" {
                    item.to_sql(context)
                } else {
                    Ok(String::default())
                }
            },
            SqlGraphEntity::Enum(item) => item.to_sql(context),
            // ...
        }
    }
}

Then the same trait goes on the types which SqlGraphEntity wraps, here's what the function looks like for enums:

impl ToSql for InventoryPostgresEnum {
    #[tracing::instrument(
        level = "debug",
        err,
        skip(self, context),
        fields(identifier = %self.rust_identifier())
    )]
    fn to_sql(&self, context: &super::PgxSql) -> eyre::Result<String> {
        let self_index = context.enums[self];
        let sql = format!(
            "\n\
                    -- {file}:{line}\n\
                    -- {full_path}\n\
                    CREATE TYPE {schema}{name} AS ENUM (\n\
                        {variants}\
                    );\
                ",
            schema = context.schema_prefix_for(&self_index),
            full_path = self.full_path,
            file = self.file,
            line = self.line,
            name = self.name,
            variants = self
                .variants
                .iter()
                .map(|variant| format!("\t'{}'", variant))
                .collect::<Vec<_>>()
                .join(",\n")
                + "\n",
        );
        tracing::debug!(%sql);
        Ok(sql)
    }
}

Other implementations, such as on functions, are a bit more complicated. These functions are tasked with determining, for example, the correct name for an argument in SQL, or the schema which contains the function.

Clockwork. - Josh Redd
Clockwork.

Closing thoughts

With the toolkit pgx provides, users of Rust are able to develop PostgreSQL extensions using familiar tooling and workflows. There's still a lot of things we'd love to refine and add to the toolkit, but we think you can start using it, like we do in production at TCDI. We think it's even fun to use.

Thanks to @dtolnay for making many of the crates discussed here, as well as being such a kind and wise person to exist in the orbit of. Also to Eric Ridge for all the PostgreSQL knowledge, and TCDI for employing me to work on this exceptionally fun stuff!

6c3fee57ca834cb9f133106444ebe3bf12a9628e