Allow plugins to return custom errors

pgdog-macros/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pgdog-macros"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2024"
 authors = ["Lev Kokotov <[email protected]>"]
 license = "MIT"

pgdog-macros/src/lib.rs

@@ -25,6 +25,14 @@ pub fn plugin(_input: TokenStream) -> TokenStream {
             }
         }
 
+        #[unsafe(no_mangle)]
+        pub unsafe extern "C" fn pgdog_plugin_lib_version(output: *mut pgdog_plugin::PdStr) {
+            let version = pgdog_plugin::comp::plugin_lib_version();
+            unsafe {
+                *output = version;
+            }
+        }
+
         #[unsafe(no_mangle)]
         pub unsafe extern "C" fn pgdog_pg_query_version(output: *mut pgdog_plugin::PdStr) {
             let version: pgdog_plugin::PdStr = option_env!("PGDOG_PGQUERY_VERSION")

pgdog-plugin/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pgdog-plugin"
-version = "0.1.8"
+version = "0.1.9"
 edition = "2021"
 license = "MIT"
 authors = ["Lev Kokotov <[email protected]>"]
@@ -18,7 +18,7 @@ libloading = "0.8"
 libc = "0.2"
 tracing = "0.1"
 pg_query = "6.1.1"
-pgdog-macros = { path = "../pgdog-macros", version = "0.1.1" }
+pgdog-macros = { path = "../pgdog-macros", version = "0.1.2" }
 toml = "0.9"
 
 [build-dependencies]

pgdog-plugin/include/types.h

@@ -61,10 +61,23 @@ typedef struct PdRouterContext {
     PdParameters params;
 } PdRouterContext;
 
+/**
+ * ErrorResponse.
+ */
+typedef struct PdErrorResponse {
+   const char *severity;
+   const char *code;
+   const char *message;
+   const char *detail;
+   const char *context;
+   const char *file;
+   const char *routine;
+} PdErrorResponse;
+
 /**
  * Routing decision returned by the plugin.
  */
- typedef struct PdRoute {
+typedef struct PdRoute {
      /** Which shard the query should go to.
       *
       * `-1` for all shards, `-2` for unknown, this setting is ignored.
@@ -75,4 +88,8 @@ typedef struct PdRouterContext {
       * `1` for `true`, `0` for `false`, `2` for unknown, this setting is ignored.
       */
      uint8_t read_write;
+     /**
+      * Specific error response to return, if blocking query.
+      */
+      PdErrorResponse error_response;
  } PdRoute;

pgdog-plugin/src/bindings.rs

@@ -308,6 +308,37 @@ const _: () = {
     ["Offset of field: PdRouterContext::params"]
         [::std::mem::offset_of!(PdRouterContext, params) - 40usize];
 };
+#[doc = " ErrorResponse."]
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct PdErrorResponse {
+    pub severity: *const ::std::os::raw::c_char,
+    pub code: *const ::std::os::raw::c_char,
+    pub message: *const ::std::os::raw::c_char,
+    pub detail: *const ::std::os::raw::c_char,
+    pub context: *const ::std::os::raw::c_char,
+    pub file: *const ::std::os::raw::c_char,
+    pub routine: *const ::std::os::raw::c_char,
+}
+#[allow(clippy::unnecessary_operation, clippy::identity_op)]
+const _: () = {
+    ["Size of PdErrorResponse"][::std::mem::size_of::<PdErrorResponse>() - 56usize];
+    ["Alignment of PdErrorResponse"][::std::mem::align_of::<PdErrorResponse>() - 8usize];
+    ["Offset of field: PdErrorResponse::severity"]
+        [::std::mem::offset_of!(PdErrorResponse, severity) - 0usize];
+    ["Offset of field: PdErrorResponse::code"]
+        [::std::mem::offset_of!(PdErrorResponse, code) - 8usize];
+    ["Offset of field: PdErrorResponse::message"]
+        [::std::mem::offset_of!(PdErrorResponse, message) - 16usize];
+    ["Offset of field: PdErrorResponse::detail"]
+        [::std::mem::offset_of!(PdErrorResponse, detail) - 24usize];
+    ["Offset of field: PdErrorResponse::context"]
+        [::std::mem::offset_of!(PdErrorResponse, context) - 32usize];
+    ["Offset of field: PdErrorResponse::file"]
+        [::std::mem::offset_of!(PdErrorResponse, file) - 40usize];
+    ["Offset of field: PdErrorResponse::routine"]
+        [::std::mem::offset_of!(PdErrorResponse, routine) - 48usize];
+};
 #[doc = " Routing decision returned by the plugin."]
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
@@ -316,11 +347,15 @@ pub struct PdRoute {
     pub shard: i64,
     #[doc = " Is the query a read and should go to a replica?\n\n `1` for `true`, `0` for `false`, `2` for unknown, this setting is ignored."]
     pub read_write: u8,
+    #[doc = " Specific error response to return, if blocking query."]
+    pub error_response: PdErrorResponse,
 }
 #[allow(clippy::unnecessary_operation, clippy::identity_op)]
 const _: () = {
-    ["Size of PdRoute"][::std::mem::size_of::<PdRoute>() - 16usize];
+    ["Size of PdRoute"][::std::mem::size_of::<PdRoute>() - 72usize];
     ["Alignment of PdRoute"][::std::mem::align_of::<PdRoute>() - 8usize];
     ["Offset of field: PdRoute::shard"][::std::mem::offset_of!(PdRoute, shard) - 0usize];
     ["Offset of field: PdRoute::read_write"][::std::mem::offset_of!(PdRoute, read_write) - 8usize];
+    ["Offset of field: PdRoute::error_response"]
+        [::std::mem::offset_of!(PdRoute, error_response) - 16usize];
 };

pgdog-plugin/src/comp.rs

@@ -6,3 +6,8 @@ use crate::PdStr;
 pub fn rustc_version() -> PdStr {
     env!("RUSTC_VERSION").into()
 }
+
+/// Version of pgdog-plugin itself.
+pub fn plugin_lib_version() -> PdStr {
+    env!("CARGO_PKG_VERSION").into()
+}

pgdog-plugin/src/context.rs

@@ -1,9 +1,10 @@
 //! Context passed to and from the plugins.
 
-use std::ops::Deref;
+use std::{fmt::Debug, ops::Deref};
 
 use crate::{
-    bindings::PdRouterContext, parameters::Parameters, PdParameters, PdRoute, PdStatement,
+    bindings::PdRouterContext, error_response::ErrorResponse, parameters::Parameters,
+    PdErrorResponse, PdParameters, PdRoute, PdStatement,
 };
 
 /// PostgreSQL statement, parsed by [`pg_query`].
@@ -449,6 +450,7 @@ impl Route {
             ffi: PdRoute {
                 shard: shard.into(),
                 read_write: read_write.into(),
+                error_response: PdErrorResponse::none(),
             },
         }
     }
@@ -462,6 +464,7 @@ impl Route {
             ffi: PdRoute {
                 shard: -2,
                 read_write: 2,
+                error_response: PdErrorResponse::none(),
             },
         }
     }
@@ -473,7 +476,48 @@ impl Route {
             ffi: PdRoute {
                 shard: -3,
                 read_write: 2,
+                error_response: PdErrorResponse::none(),
             },
         }
     }
+
+    /// Block the query from being sent to a database and return the provided error message.
+    ///
+    /// # Arguments
+    ///
+    /// * `error`: Error response to return to the client.
+    ///
+    pub fn error(error: ErrorResponse) -> Route {
+        Self {
+            ffi: PdRoute {
+                shard: -3,
+                read_write: 2,
+                error_response: error.into(),
+            },
+        }
+    }
+
+    /// Get the error set by the plugin on the route, if any.
+    pub fn get_error(&self) -> Option<ErrorResponse> {
+        if self.error_response.is_none() {
+            None
+        } else {
+            Some(ErrorResponse::from(self.ffi.error_response))
+        }
+    }
+}
+
+impl Debug for Route {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Route")
+            .field("shard", &self.ffi.shard)
+            .field("read_write", &self.ffi.read_write)
+            .finish()
+    }
+}
+
+impl Drop for Route {
+    fn drop(&mut self) {
+        unsafe { self.ffi.error_response.deallocate() };
+    }
 }