Training Hub › Rust & Cross-Language

Rust PostgreSQL Extensions with pgrx

Systems Programming advanced Building with AI ~35 min read

Overview

PostgreSQL’s extension system is one of its most powerful features, allowing developers to add custom data types, operators, index methods, and foreign data wrappers that run inside the database process. Traditionally, extensions are written in C using PostgreSQL’s internal APIs. The pgrx framework (formerly pgx) makes it possible to write these extensions in Rust, providing memory safety guarantees while still producing native shared libraries that PostgreSQL loads directly.

Rust is an ideal language for PostgreSQL extensions because extensions run in the database server’s address space—a segfault in an extension crashes the entire database. C extensions are notoriously prone to buffer overflows, use-after-free bugs, and null pointer dereferences that take down production databases. Rust’s ownership system eliminates these classes of bugs at compile time, while pgrx provides safe Rust abstractions over PostgreSQL’s C internals.

As an example, the supabase-wrappers pattern can be used to build Foreign Data Wrappers (FDWs) that let PostgreSQL query external data sources (including a search index service) as if they were local tables. The pgrx framework handles the complex FFI boundary between Rust and PostgreSQL’s C runtime, including memory context management, error handling via longjmp, and the SPI (Server Programming Interface) for executing SQL from within extensions.

Core Concepts

pg_module_magic and Extension Loading

Every PostgreSQL extension must declare a “magic block” that the server uses to verify binary compatibility. In C, this is the PG_MODULE_MAGIC macro. In pgrx, the pg_module_magic!() macro generates this block. When PostgreSQL loads the shared library (.so/.dylib), it checks the magic block against its own version. A mismatch means the extension was compiled for a different PostgreSQL major version and loading is refused.

Memory Contexts

PostgreSQL uses a hierarchical memory context system instead of malloc/free. Each memory context is a pool that can be freed all at once. The CurrentMemoryContext changes throughout query execution: per-tuple contexts are reset after each row, transaction contexts persist for the transaction duration, and TopMemoryContext lives for the entire session. pgrx integrates with this system, ensuring Rust-allocated objects are placed in the correct context and freed when the context resets.

Error Handling and longjmp

PostgreSQL uses C’s setjmp/longjmp for error handling (similar to exceptions, but without destructors). When a PostgreSQL error occurs, ereport(ERROR, ...) calls longjmp, unwinding the stack back to the nearest error handler. This skips Rust’s drop calls, potentially leaking resources. pgrx provides PgTryBuilder to catch these errors and run cleanup code, bridging PostgreSQL’s C error model with Rust’s RAII pattern.

Key Concept: #[pg_extern] Attribute

The #[pg_extern] attribute macro transforms a Rust function into a PostgreSQL-callable function. It generates the C ABI wrapper, argument unmarshalling from PostgreSQL’s Datum type, null handling, and error translation. A function annotated with #[pg_extern] can be called directly from SQL after CREATE EXTENSION.

Key Concept: Foreign Data Wrappers (FDW)

FDWs allow PostgreSQL to query external data sources through standard SQL. The FDW API defines callbacks for planning (estimating row counts and costs), scanning (iterating rows), and modification (insert/update/delete). pgrx and supabase-wrappers provide Rust traits that map to these callbacks, letting you implement an FDW by filling in a Rust trait rather than wiring up C function pointers manually.

Key Concept: SPI — Server Programming Interface

SPI allows extension code to execute SQL queries within the server process. pgrx wraps SPI in a safe Rust API: Spi::connect(|client| { client.select("SELECT ...", ...) }). SPI queries run in the same transaction as the calling function, see the same snapshot, and respect the same locks. This enables extensions to read and write tables as part of their logic.

How It Works

Basic pgrx Extension Structure

A pgrx extension starts with cargo pgrx new, which scaffolds a Cargo project with the right dependencies and build configuration. The core of the extension is a lib.rs that declares functions:

use pgrx::prelude::*;

pg_module_magic!();

/// Calculate cosine similarity between two float arrays
#[pg_extern]
fn cosine_similarity(a: Vec<f32>, b: Vec<f32>) -> f32 {
    if a.len() != b.len() {
        pgrx::error!("vectors must have equal dimensions");
    }

    let mut dot = 0.0_f32;
    let mut norm_a = 0.0_f32;
    let mut norm_b = 0.0_f32;

    for (x, y) in a.iter().zip(b.iter()) {
        dot += x * y;
        norm_a += x * x;
        norm_b += y * y;
    }

    let denom = norm_a.sqrt() * norm_b.sqrt();
    if denom == 0.0 {
        return 0.0;
    }
    dot / denom
}

/// Search the external index and return results as a set of rows
#[pg_extern]
fn exampledb_search(
    query: &str,
    limit: i32,
) -> TableIterator<'static, (name!(id, String), name!(score, f32))> {
    let results = do_exampledb_search(query, limit as usize);

    TableIterator::new(results.into_iter().map(|r| (r.id, r.score)))
}

SQL Error Code Mapping

PostgreSQL uses 5-character SQLSTATE error codes (e.g., 22003 for numeric value out of range). pgrx maps Rust errors to these codes, allowing clients to handle specific error conditions programmatically:

use pgrx::prelude::*;

#[pg_extern]
fn safe_divide(a: f64, b: f64) -> f64 {
    if b == 0.0 {
        // Maps to SQLSTATE 22012 (division by zero)
        pgrx::error!(
            PgSqlErrorCode::ERRCODE_DIVISION_BY_ZERO,
            "division by zero"
        );
    }
    a / b
}

#[pg_extern]
fn validate_dimension(dim: i32) -> i32 {
    if dim < 1 || dim > 4096 {
        // Maps to SQLSTATE 22003 (numeric value out of range)
        pgrx::error!(
            PgSqlErrorCode::ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE,
            "dimension must be between 1 and 4096, got {}",
            dim
        );
    }
    dim
}

Foreign Data Wrapper with supabase-wrappers

The supabase-wrappers crate provides a high-level Rust trait for implementing FDWs. Instead of dealing with PostgreSQL’s low-level FDW callbacks, you implement a ForeignDataWrapper trait:

use pgrx::prelude::*;
use supabase_wrappers::prelude::*;

#[wrappers_fdw(
    version = "0.1.0",
    author = "example",
    website = "https://github.com/example"
)]
pub struct ExampleDbFdw {
    endpoint: String,
    api_key: String,
    row_iter: Option<Vec<Row>>,
    current_idx: usize,
}

impl ForeignDataWrapper<ExampleDbFdwError> for ExampleDbFdw {
    fn new(
        server: ForeignServer,
        _options: &HashMap<String, String>,
    ) -> Result<Self, ExampleDbFdwError> {
        let endpoint = server
            .options
            .get("endpoint")
            .ok_or(ExampleDbFdwError::MissingOption("endpoint"))?
            .clone();

        let api_key = server
            .options
            .get("api_key")
            .ok_or(ExampleDbFdwError::MissingOption("api_key"))?
            .clone();

        Ok(Self {
            endpoint,
            api_key,
            row_iter: None,
            current_idx: 0,
        })
    }

    fn begin_scan(
        &mut self,
        quals: &[Qual],
        columns: &[Column],
        _sorts: &[Sort],
        limit: &Option<Limit>,
        _options: &HashMap<String, String>,
    ) -> Result<(), ExampleDbFdwError> {
        // Translate PostgreSQL quals to external search service query
        let query = build_query_from_quals(quals);
        let results = self.execute_search(&query, columns, limit)?;
        self.row_iter = Some(results);
        self.current_idx = 0;
        Ok(())
    }

    fn iter_scan(&mut self, row: &mut Row) -> Result<Option<()>, ExampleDbFdwError> {
        if let Some(ref rows) = self.row_iter {
            if self.current_idx < rows.len() {
                *row = rows[self.current_idx].clone();
                self.current_idx += 1;
                return Ok(Some(()));
            }
        }
        Ok(None) // End of scan
    }

    fn end_scan(&mut self) -> Result<(), ExampleDbFdwError> {
        self.row_iter = None;
        Ok(())
    }
}
Key Concept: pg_sys FFI Boundary

The pgrx::pg_sys module provides raw Rust bindings to every PostgreSQL C function and struct. When pgrx’s safe abstractions don’t cover your use case, you can drop to unsafe { pg_sys::some_function(...) }. This is the escape hatch from safety, and the same rules apply as any Rust unsafe block: you are responsible for upholding the invariants that the C code expects.

Practical Applications

Vector Search Extension

A pgrx extension can add vector similarity search to PostgreSQL, similar to pgvector but with custom distance functions optimized for a specific embedding space. The extension defines a custom vector type, operators (&lt;-> for distance), and an index access method that uses HNSW graphs for approximate nearest neighbor search. pgrx’s type system ensures the custom type is correctly marshalled between Rust and PostgreSQL’s Datum representation.

The Foreign Data Wrapper pattern lets applications query an external search service’s indices through standard SQL. A query like SELECT * FROM exampledb_docs WHERE query = 'rust pgrx' LIMIT 10 is translated into an API call to the external search engine, and results are returned as PostgreSQL rows. This allows existing SQL-based applications to integrate search without changing their data access layer.

Custom Aggregate Functions

pgrx can implement custom PostgreSQL aggregate functions in Rust. For example, a vector_avg aggregate that computes the centroid of a set of embedding vectors, or a reciprocal_rank_fusion aggregate that combines multiple ranked result lists. These aggregates run inside PostgreSQL’s query executor, operating on data without extracting it from the database.

Background Workers

pgrx supports PostgreSQL background workers—processes that run alongside the database server. A background worker can maintain a persistent connection to an external database cluster, synchronize data, or run periodic reindexing tasks. Background workers have access to PostgreSQL’s shared memory and can communicate with backend processes through shared memory queues.

Implementation Patterns

Extension Lifecycle

The standard workflow for building and testing a pgrx extension:

  1. cargo pgrx new my_extension — Scaffold the project
  2. Implement functions with #[pg_extern] in src/lib.rs
  3. cargo pgrx test — Run tests inside a temporary PostgreSQL instance
  4. cargo pgrx run — Start a PostgreSQL instance with the extension loaded
  5. cargo pgrx package — Build a distributable .so + SQL files
  6. CREATE EXTENSION my_extension; — Install in production

Testing with pgrx

pgrx provides a test framework that spins up a real PostgreSQL instance and runs tests inside it:

#[cfg(any(test, feature = "pg_test"))]
#[pg_schema]
mod tests {
    use pgrx::prelude::*;

    #[pg_test]
    fn test_cosine_similarity() {
        let result = Spi::get_one::<f32>(
            "SELECT cosine_similarity(
                ARRAY[1.0, 0.0, 0.0]::float4[],
                ARRAY[0.0, 1.0, 0.0]::float4[]
            )"
        )
        .expect("SPI failed")
        .expect("NULL result");

        assert!((result - 0.0).abs() < 1e-6, "orthogonal vectors should have 0 similarity");
    }

    #[pg_test]
    fn test_cosine_identical() {
        let result = Spi::get_one::<f32>(
            "SELECT cosine_similarity(
                ARRAY[1.0, 2.0, 3.0]::float4[],
                ARRAY[1.0, 2.0, 3.0]::float4[]
            )"
        )
        .expect("SPI failed")
        .expect("NULL result");

        assert!((result - 1.0).abs() < 1e-6, "identical vectors should have similarity 1.0");
    }

    #[pg_test(error = "vectors must have equal dimensions")]
    fn test_dimension_mismatch() {
        Spi::run(
            "SELECT cosine_similarity(
                ARRAY[1.0, 2.0]::float4[],
                ARRAY[1.0, 2.0, 3.0]::float4[]
            )"
        ).unwrap();
    }
}

Managing Extension Versions and Migrations

PostgreSQL extensions use versioned SQL migration files (my_ext--1.0--1.1.sql) to add or modify functions between versions. pgrx generates these migration files from the Rust code, tracking which functions were added, modified, or removed. The ALTER EXTENSION my_ext UPDATE TO '1.1' command applies the migration in-place without requiring DROP EXTENSION.

-- Generated migration: my_ext--1.0--1.1.sql

-- New function added in 1.1
CREATE FUNCTION exampledb_search(query text, limit_val integer)
RETURNS TABLE (id text, score real)
LANGUAGE c
AS 'MODULE_PATHNAME', 'exampledb_search_wrapper';

-- Existing function signature unchanged, but implementation updated
-- (replacing the .so is sufficient, no SQL change needed)
Key Concept: Datum and Type Safety

PostgreSQL represents all values as Datum—a pointer-sized integer that can hold a value directly (for small types like int4) or point to heap-allocated data (for variable-length types like text). pgrx’s FromDatum and IntoDatum traits convert between Rust types and Datums safely, with compile-time checks that prevent passing an i32 where PostgreSQL expects a text.

Performance Considerations

Common Pitfalls

Further Reading