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

Photo Credit Photograph

Préambule

Cet article 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:

Important: Cette série d'articles n'est pas une introduction au langage Rust ! Il existe d'autres ressources pour ça (notamment ici et là). On se focalisera ici sur la mécanique interne d'une extension Postgres.

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

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.

Démarrer

Si vous n'avez pas encore fait, installer PGRX, reportez-vous à la section "Démarrage" de l'article "Episode 1 : Un nouveau monde".

ou sinon utilisez l'image docker:

docker run -it --volume `pwd`:/pgrx daamien/pgrx

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

struct ParapetHooks {
}

impl pgrx::hooks::PgHooks for ParapetHooks {

    fn post_parse_analyze(
        &mut self,
        parse_state: PgBox<pg_sys::ParseState>,
        query: PgBox<pg_sys::Query>,
        jumble_state: Option<PgBox<pgrx::JumbleState>>,
        prev_hook: fn(
            parse_state: PgBox<pg_sys::ParseState>,
            query: PgBox<pg_sys::Query>,
            jumble_state: Option<PgBox<pgrx::JumbleState>>,
        ) -> pgrx::hooks::HookResult<()>,
    ) -> pgrx::hooks::HookResult<()> {
            debug1!("Parapet: Query hooked");
            // call the previous hook (if any)
            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.

NOTE EN PASSANT: Si vous avez déja utilisé ce genre de crochet, vous savez peut-être que la signature de ces fonctions n'est pas stable d'une version majeure à l'autre. Par exemple le paramètre jumble_state ci-dessus a été introduit dans la version 14 de PostgreSQL. Lorsque l'on écrit une extension en C, cela impose d'utiliser les infames macros #if PG_VERSION_NUM. Là encore PGRX simplifie tout et "stabilise" les signatures crochets. Ce n'est pas magique, le paramètre jumble_state reste inexploitable avant la version 14. Mais ça rend le code beaucoup plus lisible et agréable.

Il ne reste plus qu'à instancier un variable globale de type ParapetHooks et l'enregistrer (register) dans la fonction _PG_init() qui se déclenchera au moment du chargement de l'extension.

// dans le fichier src/lib.rs
static mut HOOKS: ParapetHooks = ParapetHooks {};

#[pg_guard]
pub unsafe extern "C" fn _PG_init() {
    pgrx::hooks::register_hook(&mut HOOKS);
}

Maintenant lançons l'extension et testons !

cargo pgrx run

Attention! Cette extension utilise la fonction _PG_ini() 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 :

`` sql 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` :

``` sql
\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 {
    // The "jointree" is a raw pointer
    // as we dereference that raw pointer, we take take the responsibility that
    // it may point to an incorrect location.
    // This is why we need to make this operation as `unsafe`
    let jointree = unsafe { *query.jointree };
    debug1!("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 :

    fn post_parse_analyze(
        // (les paramètres de la fonction sont omis pour plus de clarté)
    ) -> pgrx::hooks::HookResult<()> {
            debug1!("Parapet: Query hooked");
            match query.commandType {
                pg_sys::CmdType_CMD_DELETE => check_delete(&query),
                _ => ()
            }
            // call the previous hook (if any)
            prev_hook(parse_state, query, jumble_state)
    }

Et testons à nouveau !

cargo pgrx run

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

SET client_min_messages TO DEBUG1;

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:

debug1!("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 le système de sécurité des threads 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() {
    pgrx::hooks::register_hook(&mut HOOKS);

    GucRegistry::define_bool_guc(
            "pg_parapet.force_where_on_delete",
            "Activate the Parapet protection on DELETE commands",
            "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:


DELETE FROM t ;
DELETE 0

SET pg_parapet.force_where_on_delete TO true;

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

RESET pg_parapet.force_where_on_delete;

DELETE FROM t ;
DELETE 0

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;").expect("SPI Failed");
        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;").expect("SPI Failed");
        Spi::run("CREATE TABLE t AS SELECT 1;").expect("SPI Failed");
        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;").expect("SPI Failed");
        Spi::run("CREATE TABLE t AS SELECT 1;").expect("SPI Failed");
        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.

Keys	Action
`?`	Open this help
`n`	Next page
`p`	Previous page
`s`	Search