使用pgrx開發postgre外掛

GaN8373發表於2024-11-12

pg外掛開發建議在linux環境下進行, windows可以採用虛擬機器的方式開發.

安裝虛擬機器環境

  1. 準備虛擬機器, 使用的是ubuntu22.04, 太新的版本會依賴拉取有問題
  2. 安裝開發工具和依賴:
    ​sudo apt-get update sudo apt-get install -y git clang gcc make build-essential libz-dev zlib1g-dev strace libssl-dev pkg-config libreadline-dev libpg-dev flex bison libxml2-dev libxslt-dev libxml2-utils xsltproc​
  3. 安裝rust: curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh​
  4. 安裝cargo: sudo apt-get install cargo​
  5. 安裝pgrx: cargo install cargo-pgrx​
  6. 初始化pgrx: cargo pgrx init​ 這會安裝12 - 17版本的pg.
    耐心等待所有步驟完成後, 就可以使用pgrx新建模板來寫外掛了

hello pgrx

執行命令新建pg擴充套件專案: cargo pgrx new <project_name>​
使用生成的專案模板進行外掛開發

除錯

如果有需要, 可以修改自己要用的pg版本配置檔案來方便除錯, 例如要用pg16就改.pgrx/data-16

  1. pg_hba.conf和pg_ident.conf 支援外部工具連線pg
  2. postgresql.conf設定日誌輸出級別
    修改log_min_messages和client_min_messages

執行預設模板外掛的函式

進入pgrx新建的專案目錄後, 在控制檯執行cargo pgrx run pg16表示在postgre16版本執行該外掛.
預設埠起始是28800, 因此pg16開放的埠是28816(pg17就是28817),如果開放允許外部連線可以使用這個埠進去(pgrx run命令會自動接入到pg控制檯模式)
執行sql安裝pg擴充套件:

create extension <project_name>;
SELECT extname AS "Extension Name", extversion AS "Version"
FROM pg_extension;

<project_name>的來源是自己建立的專案名稱.

除錯用的宏

例: use pgrx::{debug1, debug5}
按照pg級別使用, 有:

  1. debug5!
  2. debug4!
  3. debug3!
  4. debug2!
  5. debug1!
  6. info!
  7. notice!
  8. warning!
  9. error!
  10. log!
  11. FATAL!
  12. PANIC!

自定義型別

see: 自定義型別官方示例
自定義型別需要實現序列化和反序列化:

#[derive(PostgresType, Deserialize, Serialize)]
pub struct MyType {}

過載運算子

see: 過載運算子官方示例

#[pg_operator(immutable, parallel_safe)]
#[opname(||)]
fn mytype_concact(mut left: MyType, right: MyType) -> MyType {
    left.0.push_str(&right.0);
    left
}

處理panic

官方提供了一個叫做PgTryBuilder的實現, 用來捕獲rust中的panic.

#[pg_extern]
fn maybe_panic(panic: bool, trap_it: bool, message: &str) {
    PgTryBuilder::new(|| {
        if panic {
            panic!("panic says: {}", message)
            // std::panic::panic_any(42)
        }
    })
    .catch_rust_panic(|cause| {
        // we can catch general Rust panics.  The `error` argument is a pg_sys::panic::CaughtError::RustPanic
        // enum variant with the payload from the originating panic
        if trap_it {
            if let CaughtError::RustPanic { ereport, payload } = &cause {
                warning!("{:#?}", ereport);
                if let Some(s) = payload.downcast_ref::<String>() {
                    // we have access to the panic!() message
                    warning!("{}", s);
                    return;
                } else {
                    // this won't happen with this example, but say the `panic_any(42)` was used
                    // instead.  Then we'd be here, and we can just raise another `panic!()`, which
                    // will be what's ultimately reported to Postgres.
                    //
                    // In this case, Postgres' LOCATION error slot will be this line, and the CONTEXT
                    // slot will show the line number of the original `panic_any(42)` above
                    panic!("panic payload not a `String`");
                }
            }
            unreachable!("internal error:  `CaughtError` not a `::RustPanic`");
        } else {
            cause.rethrow()
        }
    })
    // finally block always runs after the catch handlers finish (even if they rethrow or raise
    // their own panic, like in this case)
    .finally(|| warning!("FINALLY!"))
    .execute()
}

在外掛中執行SQL

see: 官方spi示例
see: hook列表(沒有試過)
執行SQL需要用到pgrx提供的spi介面
這是一個在hook中執行sql的示例:

#[pg_guard]
pub unsafe extern "C" fn _PG_init() {
        PgTryBuilder::new(|| {
            let vec = generate_rsa_key().unwrap();
            Spi::connect(move |mut client| {
                let init_schema = format!(
                    r#"
                    create schema if not exists  {};
                "#,
                    PROJECT_SCHEMA
                );

                match client.update(init_schema.as_str(), None, None) {
                    Ok(_) => {}
                    Err(_) => {
                        return;
                    }
                };

                let init_sql = format!(
                    r#"
                create table if not exists  {}.t_rsa_key(
                    id serial8 not null primary key,
                    public_key text,
                    private_key text
                );
            "#,
                    PROJECT_SCHEMA
                );

                client.update(init_sql.as_str(), None, None).unwrap();
                let insert_sql = format!(
                    r#"
            insert into {}.t_rsa_key (public_key,private_key) values ('{}','{}');
            "#,
                    PROJECT_SCHEMA, vec[0], vec[1]
                );

                client.update(insert_sql.as_str(), None, None).unwrap();
            })
        })
        .execute();
}

關於_PG_init被呼叫多次的解決方式

https://stackoverflow.com/questions/28037300/in-postgresql-why-is-pg-init-called-twice
https://github.com/pgcentralfoundation/pgrx/blob/develop/pgrx-examples/shmem/README.md
(TODO: 其實我也沒去處理這個問題)

相關文章