Rust PostgreSQL扩展开发库pgrx-pg-sys的使用，实现高性能PostgreSQL插件与系统交互

安装

在项目目录中运行以下Cargo命令：

cargo add pgrx-pg-sys

或者在Cargo.toml中添加以下行：

pgrx-pg-sys = "0.16.0"

文档

官方文档和代码仓库提供了详细的使用说明和API参考。

示例代码

以下是一个使用pgrx-pg-sys开发PostgreSQL扩展的完整示例：

// 引入必要的pgrx-pg-sys模块
use pgrx_pg_sys::{Datum, Oid, PG_FUNCTION_INFO_V1, PG_MODULE_MAGIC, Pg_magic_struct};

// 这是PostgreSQL模块必须包含的魔术结构体
PG_MODULE_MAGIC;

// 定义我们的函数信息结构
PG_FUNCTION_INFO_V1!(add_one);

// 实际的函数实现
#[no_mangle]
pub extern "C" fn add_one(fcinfo: pgrx_pg_sys::FunctionCallInfo) -> Datum {
    // 获取输入参数
    let arg = unsafe { pgrx_pg_sys::PG_GETARG_INT32(fcinfo, 0) };
    
    // 执行计算
    let result = arg + 1;
    
    // 返回结果
    unsafe { pgrx_pg_sys::PG_RETURN_INT32(result) }
}

// 另一个示例函数，计算字符串长度
PG_FUNCTION_INFO_V1!(string_length);

#[no_mangle]
pub extern "C" fn string_length(fcinfo: pgrx_pg_sys::FunctionCallInfo) -> Datum {
    // 获取文本输入参数
    let text_ptr = unsafe { pgrx_pg_sys::PG_GETARG_TEXT_PP(fcinfo, 0) };
    
    // 计算长度
    let len = unsafe { pgrx_pg_sys::VARSIZE_ANY_EXHDR(text_ptr) as i32 };
    
    // 返回结果
    unsafe { pgrx_pg_sys::PG_RETURN_INT32(len) }
}

// 更复杂的示例：聚合函数
PG_FUNCTION_INFO_V1!(int8_sum_state);
PG_FUNCTION_INFO_V1!(int8_sum_final);

#[no_mangle]
pub extern "C" fn int8_sum_state(
    fcinfo: pgrx_pg_sys::FunctionCallInfo,
) -> pgrx_pg_sys::Datum {
    // 获取状态值和输入值
    let state = unsafe { pgrx_pg_sys::PG_GETARG_INT64(fcinfo, 0) };
    let value = unsafe { pgrx_pg_sys::PG_GETARG_INT64(fcinfo, 1) };
    
    // 更新状态
    let new_state = state + value;
    
    // 返回新状态
    unsafe { pgrx_pg_sys::PG_RETURN_INT64(new_state) }
}

#[no_mangle]
pub extern "C" fn int8_sum_final(
    fcinfo: pgrx_pg_sys::FunctionCallInfo,
) -> pgrx_pg_sys::Datum {
    // 获取最终状态值
    let state = unsafe { pgrx_pg_sys::PG_GETARG_INT64(fcinfo, 0) };
    
    // 返回结果
    unsafe { pgrx_pg_sys::PG_RETURN_INT64(state) }
}

完整示例代码

以下是一个更完整的PostgreSQL扩展示例，包含多种函数类型和错误处理：

//! 一个完整的PostgreSQL扩展示例

use pgrx_pg_sys::{
    Datum, Oid, PG_FUNCTION_INFO_V1, PG_MODULE_MAGIC, Pg_magic_struct, 
    ERROR, ereport, WARNING, ERRCODE_INVALID_PARAMETER_VALUE
};

// PostgreSQL模块魔术结构体
PG_MODULE_MAGIC;

// 1. 简单标量函数示例：计算平方
PG_FUNCTION_INFO_V1!(square);

#[no_mangle]
pub extern "C" fn square(fcinfo: pgrx_pg_sys::FunctionCallInfo) -> Datum {
    // 检查参数个数
    unsafe {
        if pgrx_pg_sys::PG_NARGS() != 1 {
            ereport(
                ERROR,
                pgrx_pg_sys::errcode(ERRCODE_INVALID_PARAMETER_VALUE),
                pgrx_pg_sys::errmsg("square函数需要一个参数")
            );
        }
    }

    // 获取输入参数
    let arg = unsafe { pgrx_pg_sys::PG_GETARG_FLOAT8(fcinfo, 0) };
    
    // 计算平方
    let result = arg * arg;
    
    // 返回结果
    unsafe { pgrx_pg_sys::PG_RETURN_FLOAT8(result) }
}

// 2. 文本处理函数示例：字符串反转
PG_FUNCTION_INFO_V1!(reverse_string);

#[no_mangle]
pub extern "C" fn reverse_string(fcinfo: pgrx_pg_sys::FunctionCallInfo) -> Datum {
    // 获取文本输入参数
    let text_ptr = unsafe { pgrx_pg_sys::PG_GETARG_TEXT_PP(fcinfo, 0) };
    
    // 获取文本长度和数据指针
    let len = unsafe { pgrx_pg_sys::VARSIZE_ANY_EXHDR(text_ptr) };
    let data_ptr = unsafe { pgrx_pg_sys::VARDATA_ANY(text_ptr) };
    
    // 将文本转换为Rust字符串切片
    let text_slice = unsafe { std::slice::from_raw_parts(data_ptr as *const u8, len) };
    
    // 反转字符串
    let mut reversed = text_slice.to_vec();
    reversed.reverse();
    
    // 创建新的文本值
    let result = unsafe {
        let result_ptr = pgrx_pg_sys::palloc(len + pgrx_pg_sys::VARHDRSZ) as *mut pgrx_pg_sys::text;
        pgrx_pg_sys::SET_VARSIZE(result_ptr, len + pgrx_pg_sys::VARHDRSZ);
        std::ptr::copy_nonoverlapping(
            reversed.as_ptr(),
            pgrx_pg_sys::VARDATA(result_ptr),
            len
        );
        result_ptr
    };
    
    // 返回结果
    unsafe { pgrx_pg_sys::PG_RETURN_TEXT_P(result) }
}

// 3. 聚合函数示例：计算平均值
PG_FUNCTION_INFO_V1!(avg_state);
PG_FUNCTION_INFO_V1!(avg_final);

// 聚合状态结构体
#[repr(C)]
struct AvgState {
    sum: f64,
    count: i64,
}

#[no_mangle]
pub extern "C" fn avg_state(
    fcinfo: pgrx_pg_sys::FunctionCallInfo,
) -> pgrx_pg_sys::Datum {
    // 获取状态指针
    let state_ptr = if unsafe { pgrx_pg_sys::PG_ARGISNULL(0) } {
        // 初始化新状态
        let state_ptr = unsafe {
            pgrx_pg_sys::palloc(std::mem::size_of::<AvgState>()) as *mut AvgState
        };
        unsafe {
            (*state_ptr).sum = 0.0;
            (*state_ptr).count = 0;
        }
        state_ptr
    } else {
        unsafe { pgrx_pg_sys::PG_GETARG_POINTER(0) as *mut AvgState }
    };

    // 获取输入值
    if unsafe { !pgrx_pg_sys::PG_ARGISNULL(1) } {
        let value = unsafe { pgrx_pg_sys::PG_GETARG_FLOAT8(1) };
        unsafe {
            (*state_ptr).sum += value;
            (*state_ptr).count += 1;
        }
    }

    // 返回状态
    unsafe { pgrx_pg_sys::PG_RETURN_POINTER(state_ptr as *mut std::ffi::c_void) }
}

#[no_mangle]
pub extern "C" fn avg_final(
    fcinfo: pgrx_pg_sys::FunctionCallInfo,
) -> pgrx_pg_sys::Datum {
    // 获取状态
    let state_ptr = unsafe { pgrx_pg_sys::PG_GETARG_POINTER(0) as *mut AvgState };
    
    // 计算平均值
    let result = if unsafe { (*state_ptr).count } > 0 {
        unsafe { (*state_ptr).sum / (*state_ptr).count as f64 }
    } else {
        // 如果没有输入，返回NULL
        unsafe { pgrx_pg_sys::PG_RETURN_NULL() }
    };
    
    // 返回结果
    unsafe { pgrx_pg_sys::PG_RETURN_FLOAT8(result) }
}

// 4. 触发函数示例
PG_FUNCTION_INFO_V1!(log_trigger);

#[no_mangle]
pub extern "C" fn log_trigger(fcinfo: pgrx_pg_sys::FunctionCallInfo) -> Datum {
    unsafe {
        // 获取触发器数据
        let trigdata = pgrx_pg_sys::PG_GETARG_TRIGGER(fcinfo);
        
        // 记录日志消息
        ereport(
            WARNING,
            pgrx_pg_sys::errmsg("触发器 %s 在表 %s 上执行",
                (*trigdata).tgname,
                pgrx_pg_sys::get_rel_name((*trigdata).tgrelid)
            )
        );
        
        // 返回触发器结果
        pgrx_pg_sys::PointerGetDatum(std::ptr::null_mut())
    }
}

使用说明

1. 首先需要安装PostgreSQL开发头文件和pgrx工具链
2. 使用cargo pgrx init初始化项目
3. 编写扩展代码，如上面的示例
4. 使用cargo pgrx run测试扩展
5. 使用cargo pgrx package打包扩展

注意事项

• 使用pgrx-pg-sys需要处理大量unsafe代码
• 必须遵循PostgreSQL扩展开发的规范
• 内存管理需要特别注意，避免内存泄漏
• 错误处理需要符合PostgreSQL的机制

这个库提供了与PostgreSQL内部API的直接绑定，适合需要高性能或访问PostgreSQL内部功能的扩展开发。