Adding a Scripting Engine to a Rust CLI with Rhai

diesel-guard is a linter for Postgres migrations. It catches operations that lock tables or cause downtime before they reach production. It ships with 28 built-in checks, but over time, users started asking for custom checks:

"Our DBA requires all indexes to follow the naming convention idx_<table>_<column>. Can I enforce that?"

"Our team convention is that every new table must have an updated_at column. Can diesel-guard catch tables that don't?"

Each request is completely reasonable. But each one is unique enough that it did not make sense to add a built-in check for it. So telling users "open a PR" for a one-off team convention felt like the wrong answer.

The solution was to let users write the rule themselves, as a script, loaded at runtime, without recompiling. This post covers how I did it with Rhai.

Why Rhai

Rhai is an embedded scripting language written in Rust. You drop it into your project as you would any other crate, and scripts run within your process at runtime.

I came across this talk about it, found it interesting enough to give it a shot.

It turned out to work well for this use case. Add it to your Cargo.toml with two features:

rhai = { version = "1", features = ["serde", "sync"] }

sync makes the Rhai engine Send + Sync, which is required to store it behind a trait object. serde lets you serialize any Rust struct into a value that scripts can navigate. Both matter here, and you will see why shortly.

If you are evaluating Rhai for something performance-sensitive, the HN thread has a good discussion of the trade-offs.

The Architecture: One Trait, One Registry

Every check in diesel-guard, built-in or custom, implements the same trait:

use rhai::{AST, Dynamic, Engine};
use std::sync::Arc;

pub trait Check: Send + Sync {
    fn name(&self) -> &'static str;
    fn check(&self, node: &NodeEnum, config: &Config, ctx: &MigrationContext) -> Vec<Violation>;
}

NodeEnum is a pg_query protobuf enum with one variant per SQL statement type (IndexStmt, AlterTableStmt, CreateStmt, and so on). The Registry holds a Vec<Box<dyn Check>> and calls every check against every parsed statement.

Custom Rhai checks implement the same trait through CustomCheck:

pub struct CustomCheck {
    name: &'static str,
    engine: Arc<Engine>,
    ast: AST,
}

At check time, the registry has no way to tell whether it is calling a built-in Rust check or a Rhai script. They look the same from the outside.

To give you a sense of what this enables, here is the naming convention check from the intro. The whole thing is ten lines of Rhai:

let stmt = node.IndexStmt;
if stmt == () { return; }

let name = stmt.idxname;
if name == "" { return; }
if name.starts_with("idx_") { return; }

#{
    operation: "Index naming violation: " + name,
    problem: "Index '" + name + "' does not follow naming convention. Names must start with 'idx_'.",
    safe_alternative: "Rename: CREATE INDEX idx_" + name + " ON ...;"
}

Drop that file in a directory, point custom_checks_dir at it in diesel-guard.toml, and the check runs alongside every built-in check. How that works is what the rest of this post explains.

Bridging Typed Rust and Dynamic Scripts

This is the part that surprised me with how little code it actually needed.

pg_query gives us deeply nested protobuf types: Rust enums with dozens of variants, each with its own struct fields. Rhai scripts cannot import Rust types. The bridge is rhai::serde::to_dynamic(). Enable Rhai's serde feature, and it serializes any Serialize type into a Dynamic value that scripts can navigate like a regular map.

fn check(&self, node: &NodeEnum, config: &Config, ctx: &MigrationContext) -> Vec<Violation> {
    let dynamic_node = rhai::serde::to_dynamic(node)?;
    let dynamic_config = rhai::serde::to_dynamic(config)?;
    let dynamic_ctx = rhai::serde::to_dynamic(ctx).unwrap();

    let mut scope = rhai::Scope::new();
    scope.push("node", dynamic_node);
    scope.push("config", dynamic_config);
    scope.push("ctx", dynamic_ctx);

    self.engine.eval_ast_with_scope::<Dynamic>(&mut scope, &self.ast)
}

On the script side, protobuf enum variants become map keys and struct fields become nested map keys. A NodeEnum::IndexStmt(stmt) becomes:

{ "IndexStmt": { "concurrent": false, "idxname": "idx_users_email", ... } }

So the guard pattern in scripts looks like this:

let stmt = node.IndexStmt;
if stmt == () { return; }  // wrong node type, skip

When the node is an IndexStmt, node.IndexStmt returns the inner map. When it is anything else, it returns (), and the script exits early.

To see exactly what a statement produces, there is a built-in subcommand:

diesel-guard dump-ast --sql "CREATE INDEX idx_users_email ON users(email);"

[
  {
    "IndexStmt": {
      "concurrent": false,
      "idxname": "idx_users_email",
      "index_params": [
        { "node": { "IndexElem": { "name": "email" } } }
      ],
      "relation": { "relname": "users" },
      ...
    }
  }
]

The output strips the RawStmt and Node wrappers. What you see is what the script receives as node.

Sandboxing: Resource Limits Are Enough

Scripts run inside a sandboxed engine with four limits:

fn create_engine() -> Engine {
    let mut engine = Engine::new();
    engine.set_max_operations(100_000);  // halt infinite loops
    engine.set_max_string_size(10_000);  // prevent memory leaks
    engine.set_max_array_size(1_000);
    engine.set_max_map_size(1_000);
    engine.register_static_module("pg", create_pg_constants_module().into());
    engine
}

When max_operations is hit, Rhai returns an ErrorTerminated. We catch that and treat it as "no violation", so the linter continues normally:

Err(e) => {
    let err_str = e.to_string();
    if !err_str.contains("ErrorTerminated") {
        eprintln!("Warning: custom check '{}': runtime error: {e}", self.name);
    }
    vec![]
}

For a CLI tool checking local files written by your own team, this is sufficient. You are not running untrusted scripts from the internet.

Each invocation gets a fresh Scope, so no state leaks between migrations or between scripts.

A Stable Script API with the pg:: Module

pg_query protobuf enums are Rust types. Scripts cannot import them. If you want a script to check whether a statement is dropping an index versus dropping a table, it needs to compare stmt.remove_type against something, but that something is an integer that lives in a Rust enum.

The fix is to expose the values scripts need as a static module, registered on the engine at startup:

fn create_pg_constants_module() -> rhai::Module {
    use pg_query::protobuf::{AlterTableType, ConstrType, DropBehavior, ObjectType};

    let mut m = rhai::Module::new();
    m.set_var("OBJECT_INDEX", ObjectType::ObjectIndex as i64);
    m.set_var("OBJECT_TABLE", ObjectType::ObjectTable as i64);
    m.set_var("AT_ADD_COLUMN", AlterTableType::AtAddColumn as i64);
    m.set_var("AT_DROP_COLUMN", AlterTableType::AtDropColumn as i64);
    m.set_var("CONSTR_PRIMARY", ConstrType::ConstrPrimary as i64);
    // ...
    m
}

Scripts access them as pg::OBJECT_INDEX, pg::AT_DROP_COLUMN, and so on. If pg_query ever changes an enum value, there is one place to update it.

Return Protocol: Three Valid Shapes

Scripts return one of three things:

• () for no violation
• A map #{ operation: "...", problem: "...", safe_alternative: "..." } for one violation
• An array of those maps for multiple violations

Validation is strict. A script with a typo in a key gets a helpful error violation rather than a silent failure:

SCRIPT ERROR: my_check
Custom check returned an invalid map: 'safe_alternative' is missing
Fix the custom check script to return all three required string keys.

Compilation errors at load time are non-fatal. The function signature makes this explicit:

pub fn load_custom_checks(
    dir: &Utf8Path,
    config: &Config,
) -> (Vec<Box<dyn Check>>, Vec<ScriptError>)

Broken scripts produce warnings to stderr. Valid scripts still load and run.

A Complete Script

Here is require_concurrent_index.rhai, one of the examples included in the repo. It covers node access, context usage, conditional messages, and both the single and multi-violation return shapes:

// Require CONCURRENTLY on CREATE INDEX.
// Also flags CONCURRENTLY inside a transaction (PostgreSQL will error at runtime).
// Inspect: diesel-guard dump-ast --sql "CREATE INDEX idx ON t(id);"

let stmt = node.IndexStmt;
if stmt == () { return; }

if !stmt.concurrent {
    let idx_name = if stmt.idxname != "" { stmt.idxname } else { "(unnamed)" };
    return #{
        operation: "INDEX without CONCURRENTLY: " + idx_name,
        problem: "Creating index '" + idx_name + "' without CONCURRENTLY blocks writes on the table.",
        safe_alternative: "Use CREATE INDEX CONCURRENTLY:\n  CREATE INDEX CONCURRENTLY " + idx_name + " ON ...;"
    };
}

// CONCURRENTLY cannot run inside a transaction block
if ctx.run_in_transaction {
    let idx_name = if stmt.idxname != "" { stmt.idxname } else { "(unnamed)" };
    let hint = if ctx.no_transaction_hint != "" {
        ctx.no_transaction_hint
    } else {
        "Run this migration outside a transaction block."
    };
    #{
        operation: "INDEX CONCURRENTLY inside a transaction: " + idx_name,
        problem: "CREATE INDEX CONCURRENTLY cannot run inside a transaction block. PostgreSQL will raise an error at runtime.",
        safe_alternative: hint
    }
}

A few things worth noticing. ctx.run_in_transaction is a bool passed from the migration runner, useful for catching CONCURRENTLY inside a transaction block, which Postgres rejects at runtime. ctx.no_transaction_hint carries framework-specific instructions for how to opt out of transactions in Diesel versus SQLx. Scripts can return early with an explicit return #{ ... } or let the last expression be the return value. Integers auto-coerce to strings in concatenation.

---

If this was useful and you work with Postgres, a star on diesel-guard goes a long way. It helps other teams find the tool before a migration causes an outage.