Skip to content

Créer une extension Postgres en Rust, Episode 4 : Accrocher le parapet

Photo Credit Photograph

Préambule

Ce tutoriel fait partie d'une série consacrée au framework PGRX, un environnement de développement qui facilite la conception d'extensions PostgreSQL avec le langage Rust.

Retrouvez les autres articles de la série ci-dessous:

Chaque épisode est conçu en 2 parties : une démo de 30 minutes environ, suivi d'exercices pour aller plus loin. Les exercices sont optionnels.

Avant de démarrer ce tutorial, vérifier que PGRX est bien installé et fonctionnel:

cargo pgrx --version
cargo pgrx status

Si ce n'est pas le cas, revenir à la section installation classique ou installation via docker.

Objectifs

Dans cette 4eme partie, nous avons pour objectifs de:

  • Utiliser le mécanisme de hook de Postgres
  • Déclarer une variable GUC
  • Appeler des fonctions internes de Postgres

Contexte

Nous rentrons maintenant dans le coeur d'une fonctionnalité qui rend les extensions PostgreSQL si puissante: le système de crochets (hooks) !

Pour cet exemple, nous allons créer une extension pg_parapet qui agira comme un garde-fous afin de protéger les données d'une erreur de manipulation. La première protection consistera à interdire les commandes DELETE si elle ne possède de clause WHERE.

Créer l'extension

On peut maintenant créer une extension nommée pg_parapet avec:

cargo pgrx new pg_parapet
cd pg_parapet

Cette extension n'a pas de dépendances externes (mis à part PGRX).

Implémenter le crochet post_parse_analyze

On commence par créer une structure ParapetHooks.

On implémente le trait pgHooks pour cette structuture, ce qui nous permet de définir des fonctions crochets. Dans ce cas précis, on définit le crochet post_parse_analyze_hook qui comme son nom l'indique se déclenche à la fin de l'analyse du parser de requête.

// dans le fichier src/lib.rs

use pgrx::pg_sys;
use pgrx::pg_sys::ffi::pg_guard_ffi_boundary;

pub unsafe fn register_hooks() {
    static mut PREV_POST_PARSE_ANALYZE_HOOK: pg_sys::post_parse_analyze_hook_type = None;
    unsafe {
        PREV_POST_PARSE_ANALYZE_HOOK = pg_sys::post_parse_analyze_hook;
        pg_sys::post_parse_analyze_hook = Some(post_parse_analyze_hook);
    }

    #[pg_guard]
    unsafe extern "C-unwind" fn post_parse_analyze_hook(
        parse_state: *mut pg_sys::ParseState,
        query: *mut pg_sys::Query,
        jumble_state: *mut pg_sys::JumbleState,
    ) {
        debug1!("Parapet: Query hooked");
        // call the previous hook (if any)
        unsafe {
            if let Some(prev_hook) = PREV_POST_PARSE_ANALYZE_HOOK {
                pg_guard_ffi_boundary(|| prev_hook(parse_state, query, jumble_state));
            }
        }
    }
}

Pour l'instant notre crochet se contente d'émettre un message de DEBUG. On voit que la signature de la fonction rust correspond à la signature de la fonction C.

Il ne reste plus qu'à enregistrer (register) le crocher dans la fonction _PG_init() qui se déclenchera au moment du chargement de l'extension.

// dans le fichier src/lib.rs
#[pg_guard]
pub unsafe extern "C-unwind" fn _PG_init() {
    unsafe { register_hooks() };
}

Maintenant lançons l'extension et testons !

cargo pgrx run

Attention! Cette extension utilise la fonction _PG_init() ce qui signifie qu'elle doit être "charger la librairie" (avec la commande LOAD) et pas simplement "créer l'extension" (avec la commande CREATE EXTENSION).

Pour faciliter les tests, on charge l'extension une fois pour toute via session_preload_libraries :

ALTER DATABASE pg_parapet SET session_preload_libraries = 'pg_parapet';

Cela signifie que pour toutes les prochaines sessions ouvertes sur la base pg_parapet la librairie pg_parapet sera chargée. Il faut donc créer une nouvelle session, cela se fait simplement avec la commande \connect :

\connect
You are now connected to database "pg_parapet" as user "postgres".

SET client_min_messages TO DEBUG1;

SELECT 1;
DEBUG:  Parapet: Query hooked
 ?column?
----------
        1
(1 row)

Où est le WHERE ?

Maintenant que l'on peut intercepter les requêtes, nous allons écrire une fonction qui vérifie qu'une requête possède une clause WHERE. Le crocket post_parse_analyse prend en entrée une variable query qui contient la requête "parsée". On peut alors regarde « l'arborescence » (join tree) et regarder si elle a des conditions ( quals ). S'il n'y a pas de quals, c'est qu'il n'y pas de clause WHERE !

Ce qui nous donne:

fn has_where(query: &PgBox<pg_sys::Query>) -> bool {
    // SAFETY: The "jointree" of a Query contains the elements of
    // the FROM list and of the WHERE conditions ("quals")
    // in a DELETE Query, the joinlist is always defined
    let jointree = unsafe { PgBox::<pg_sys::FromExpr>::from_pg(query.jointree) };
    debug3!("Parapet: jointree = {:#?}", jointree);
    if jointree.quals.is_null() {
        return false;
    }
    true
}

Maintenant on déclare une fonction check_delete qui renvoie une erreur si la clause WHERE est absente:

fn check_delete(query: &PgBox<pg_sys::Query>) {
    if ! has_where(query) {
        panic!("Parapet: DELETE queries must have a WHERE clause");
    }
}

Et enfin appelons check_delete dans le crochet post_parse_analyze :

        debug1!("Parapet: Query hooked");

        // SAFETY: At this stage of the process, the query is always properly defined
        let query_boxed = unsafe { PgBox::<pg_sys::Query>::from_pg(query) };

        match query_boxed.commandType {
            pg_sys::CmdType::CMD_DELETE => check_delete(&query_boxed),
            _ => (),
        }

        // call the previous hook (if any)

Et testons à nouveau !

cargo pgrx run

Nous allons créer une table pour l'occasion:

SET client_min_messages TO DEBUG3;

CREATE TABLE t AS SELECT 1 as i;
DEBUG:  Parapet: Query hooked

SELECT * FROM t;
DEBUG:  Parapet: Query hooked
 i
---
 1
(1 row)

DELETE FROM t;
DEBUG:  Parapet: Query hooked
DEBUG:  Parapet: jointree = FromExpr {
    type_: T_FromExpr,
    fromlist: 0x000056145e43d360,
    quals: 0x0000000000000000,
}
ERROR:  Parapet: DELETE queries must have a WHERE clause

DELETE FROM t WHERE i=1;
DEBUG:  Parapet: Query hooked
DEBUG:  Parapet: jointree = FromExpr {
    type_: T_FromExpr,
    fromlist: 0x000056145e43d558,
    quals: 0x000056145e43d738,
}
DELETE 1

Le garde-fous est en place ! Plus personne ne peut utiliser la commande DELETE sans clause WHERE.

NOTE IMPORTANTE : on voit ci-dessus à quel point il est simple de d'utiliser les fonctions et les objets internes de Postgres. Cerise sur le gateau, la macro debug1! permet d'afficher le contenu de n'importe quel objet interne, ci-dessus la variable jointree est exposée avec une facilité déconcertante grace à ligne ci-dessous:

debug3!("Parapet: jointree = {:#?}", jointree);

NOTE IMPORTANTE 2 : Lors du l'appel à DELETE sans WHERE, la fonction check_delete a interrompu le traitement de la requête en lançant brutalement la macro panic!, ce qui a immédiatement tué le processus en cours...

panic!("Parapet: DELETE queries must have a WHERE clause");

Le signal PANIC est la pire chose qui puisse arriver à un processus. Cet état est généralement causé par un bug ou un comportement imprévu.

Dans le cas d'une extension écrite en C: si une extension panique, alors PostgreSQL panique également et c'est l'instance entière qui s'arrête avec une erreur de segmentation (segfault).

Ici pourtant tout s'est bien passé... L'extension pg_parapet a paniqué mais l'instance Postgres a traité cet événement comme une simple erreur et la transaction a été annulée proprement...

Comment est-ce possible ? Que s'est-il passé ?

Il s'agit pour moi de la fonctionnalité la plus bluffante de PGRX. En s'appuyant sur l'API catch_unwind de Rust, PGRX transforme toutes les erreurs, même les erreurs imprévues, en simples erreurs SQL et en provoque uniquement un ROLLBACK de la transactions en cours.

Pour le dire autrement, contrairement aux extensions écrites en C, les extensions écrites en Rust ne peuvent pas faire planter l'instance PostgreSQL.

Cette garantie de sécurité est fondamentale et c'est en soit un argument suffisant pour préférer Rust, sans même parler de perfomances, de l'écosystème de librairies disponible, ou du confort de développement.

Définir une variable GUC

Reprenons maintenant ! On l'a vu dans l'exemple plus haut, le crochet de protection est actif par défaut et il n'y a pas moyen de désactiver cette protection.

Pour rendre les choses plus configurable, il nous faudrait une variable dans le fichier postgresql.conf afin de décider si le garde-fous doit être enclenché ou pas.

Appelons cette variable pg_parapet.force_where_on_delete

On va donc déclarer une variable dans la fonction _PG_init

use pgrx::guc::*;

static GUC_FORCE_WHERE_ON_DELETE: GucSetting<bool> = GucSetting::<bool>::new(false);

pub unsafe extern "C" fn _PG_init() {

    // [...]

    GucRegistry::define_bool_guc(
            c"pg_parapet.force_where_on_delete",
            c"Activate the Parapet protection on DELETE commands",
            c"DELETE without a WHERE clause will be forbidden",
            &GUC_FORCE_WHERE_ON_DELETE,
            GucContext::Suset,
            GucFlags::default(),
    );
}

et modifions la fonction check_delete() :

fn check_delete(query: &PgBox<pg_sys::Query>) {
    if GUC_FORCE_WHERE_ON_DELETE.get() && ! has_where(query) {
        panic!("Parapet: DELETE queries must have a WHERE clause");
    }
}

Voyons ce que ça donne :

cargo pgrx run

Désormais il faut activer explicitement le paramètre GUC pour que la protection soit enclenchée:

BEGIN;

DELETE FROM t ;
DELETE 1

ROLLBACK;

L'extension est chargée mais elle est inactive....

BEGIN;

SET pg_parapet.force_where_on_delete TO true;

DELETE FROM t ;
ERROR:  Parapet: DELETE queries must have a WHERE clause

Cette fois cela a fonctionné. La transaction est en erreur mais l'instance est toujours en vie !

RESET pg_parapet.force_where_on_delete;

DELETE FROM t ;
DELETE 0

L'extension est a nouveau désactivée.

Ecrire les tests

On peut enfin tester les différentes situations avec les 3 tests suivants:

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


    #[pg_test]
    unsafe fn test_delete() {
        Spi::run("CREATE TABLE t AS SELECT 1;").unwrap();
        Spi::run("DELETE FROM t").unwrap();
    }

    #[pg_test]
    unsafe fn test_delete_with_where_is_accepted() {
        Spi::run("SET pg_parapet.force_where_on_delete to true;").unwrap();
        Spi::run("CREATE TABLE t AS SELECT 1;").unwrap();
        Spi::run("DELETE FROM t WHERE 1=1;").unwrap();
    }

    #[pg_test]
    #[should_panic]
    unsafe fn test_delete_without_where_is_rejected() {
        Spi::run("SET pg_parapet.force_where_on_delete to true;").unwrap();
        Spi::run("CREATE TABLE t AS SELECT 1;").unwrap();
        Spi::run("DELETE FROM t").unwrap();
    }

}

Pour aller plus loin...

Voici quelques pistes pour continuer dans cette voie

  • Ajouter une protection similaire pour la commande UPDATE avec une variable pg_parapet.force_where_on_update.

  • Ajouter une protection pour que seuls les super-utilisateurs puisse utiliser la commande TRUNCATE. Pour cela, il faudra utiliser le crochet process_utility_hook, vérifier que la commande (utilityStmt) est de type pg_sys::NodeTag::T_TruncateStmt et enfin vérifier les privilèges de l'utilisateur avec pg_sys::superuser()

Un exemple d'implémentation est disponible sur:

https://gitlab.com/daamien/pgrx-tuto/-/blob/main/pg_parapet/src/lib.rs

Ecrire la même chose en C ?

En guise de comparaison, on pourra regarder l'extension pg-safeupdate qui fait sensiblement la même chose (mais uniquement pour PG14+).

Bilan

On a vu dans cet épisode que PGRX donne un accès direct aux crochets et aux objets internes de Postgres. Le périmètre fonctionnel est comparable à celui des extensions écrites en C mais avec une fiablité renforcée.