feat: plugins can now return sharding keys

.github/workflows/ci-new-parser.yml

@@ -76,6 +76,9 @@ jobs:
     env:
       PGDOG_BIN: ${{ github.workspace }}/target/debug/pgdog
       PGDOG_PLUGIN_FEATURES: new_parser
+      # Run only the plugin suite under gdb so crashes include a full backtrace
+      # in integration/log.txt, which stop_pgdog prints on job teardown.
+      PGDOG_GDB: ${{ matrix.name == 'plugins' && '1' || '0' }}
     steps:
       - uses: actions/checkout@v6
       - name: Install CI deps

.github/workflows/ci.yml

@@ -97,6 +97,9 @@ jobs:
       # All suites run against the llvm-cov instrumented debug binary so we
       # get line coverage from integration tests, not just unit tests.
       PGDOG_BIN: ${{ github.workspace }}/target/llvm-cov-target/debug/pgdog
+      # Run only the plugin suite under gdb so crashes include a full backtrace
+      # in integration/log.txt, which stop_pgdog prints on job teardown.
+      PGDOG_GDB: ${{ matrix.name == 'plugins' && '1' || '0' }}
       # NB: cargo llvm-cov report only globs profraw files at the top of
       # target/llvm-cov-target, so LLVM_PROFILE_FILE must point there.
       LLVM_PROFILE_FILE: ${{ github.workspace }}/target/llvm-cov-target/${{ matrix.name }}-%p-%m.profraw

Cargo.toml

@@ -19,7 +19,7 @@ members = [
 edition = "2024"
 
 [workspace.dependencies]
-pgdog-plugin = { path = "./pgdog-plugin", version = "0.4.0", default-features = false }
+pgdog-plugin = { path = "./pgdog-plugin", version = "0.5.0", default-features = false }
 pgdog-config = { path = "./pgdog-config", version = "0.1.0" }
 pgdog-postgres-types = { path = "./pgdog-postgres-types"}
 pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "4843f8bb01c1b7f844d2b8c43de551d3a493d333" }

integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs

@@ -2,7 +2,7 @@
 //! The plugin uses the same version of pgdog-plugin as PgDog and the same rustc version,
 //! so it should be loaded and executed.
 
-use pgdog_plugin::{plugin, Context, PdStr, Plugin, Route};
+use pgdog_plugin::{plugin, Context, PdStr, Plugin, ReadWrite, Route};
 use std::sync::OnceLock;
 
 plugin!(TestPlugin);
@@ -43,6 +43,8 @@ impl Plugin for TestPlugin {
             std::fs::write(&file_path, "route method was called").unwrap();
         });
 
-        Route::unknown()
+        // Allocate a sharding key in the plugin. PgDog takes ownership of the
+        // returned route and deallocates this string with it.
+        Route::with_sharding_keys(vec!["plugin-owned-key".into()], ReadWrite::Unknown)
     }
 }

pgdog-plugin/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pgdog-plugin"
-version = "0.4.0"
+version = "0.5.0"
 edition = "2024"
 license = "MIT"
 authors = ["Lev Kokotov <lev.kokotov@gmail.com>"]
@@ -18,6 +18,7 @@ libloading = "0.8"
 pg_query = { git = "https://github.com/pgdogdev/pg_query.rs.git", rev = "97019d0c13ad0b888fe91ee5bed5448b5f409cdd", optional = true }
 pg_raw_parse = { workspace = true, optional = true }
 pgdog-postgres-types.workspace = true
+uuid.workspace = true
 
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }

pgdog-plugin/src/context.rs

@@ -1,6 +1,9 @@
 //! Context passed to and from the plugins.
 
-use crate::parameters::Parameters;
+use crate::{
+    parameters::Parameters,
+    value::{ShardingKeys, Value},
+};
 
 /// Context information provided by PgDog to the plugin at statement execution. It contains the actual statement and several metadata about
 /// the state of the database cluster:
@@ -355,6 +358,13 @@ impl From<ReadWrite> for u8 {
 /// // No routing information is available. PgDog will ignore it
 /// // and make its own decision.
 /// let route = Route::unknown();
+///
+/// // Let PgDog calculate shards for keys extracted by the plugin.
+/// let route = Route::with_sharding_keys(
+///     vec![123_i64.into(), "tenant-a".into()],
+///     ReadWrite::Read,
+/// );
+/// ```
 #[repr(C)]
 pub struct Route {
     /// Which shard the query should go to.
@@ -365,6 +375,18 @@ pub struct Route {
     ///
     /// `1` for `true`, `0` for `false`, `2` for unknown, this setting is ignored.
     pub read_write: u8,
+    /// Plugin-allocated sharding keys.
+    sharding_keys: ShardingKeys,
+}
+
+impl std::fmt::Debug for Route {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Route")
+            .field("shard", &self.shard)
+            .field("read_write", &self.read_write)
+            .field("sharding_keys", &self.sharding_keys())
+            .finish()
+    }
 }
 
 impl Default for Route {
@@ -385,26 +407,70 @@ impl Route {
         Self {
             shard: shard.into(),
             read_write: read_write.into(),
+            sharding_keys: ShardingKeys::default(),
+        }
+    }
+
+    /// Create a route from sharding keys extracted from the query.
+    ///
+    /// The plugin transfers ownership of the values to PgDog when it returns
+    /// this route. PgDog can then pass each typed value to the sharding
+    /// function selected for the query.
+    pub fn with_sharding_keys(sharding_keys: Vec<Value>, read_write: ReadWrite) -> Route {
+        Self {
+            shard: Shard::Unknown.into(),
+            read_write: read_write.into(),
+            sharding_keys: ShardingKeys::new(sharding_keys),
         }
     }
 
+    /// Sharding keys extracted by the plugin.
+    pub fn sharding_keys(&self) -> &ShardingKeys {
+        &self.sharding_keys
+    }
+
     /// Create new route with no sharding or read/write information.
     /// Use this if you don't want your plugin to do query routing.
     /// Plugins that do something else with queries, e.g., logging, metrics,
     /// can return this route.
     pub fn unknown() -> Route {
-        Self {
-            shard: -2,
-            read_write: 2,
-        }
+        Self::new(Shard::Unknown, ReadWrite::Unknown)
     }
 
     /// Block the query from being sent to a database. PgDog will abort the query
     /// and return an error to the client, telling them which plugin blocked it.
     pub fn block() -> Route {
-        Self {
-            shard: -3,
-            read_write: 2,
-        }
+        Self::new(Shard::Blocked, ReadWrite::Unknown)
+    }
+}
+
+#[cfg(test)]
+mod route_tests {
+    use super::*;
+
+    #[test]
+    fn route_owns_sharding_keys() {
+        let uuid = uuid::Uuid::nil();
+        let route = Route::with_sharding_keys(
+            vec![1_i64.into(), "tenant".into(), uuid.into()],
+            ReadWrite::Read,
+        );
+
+        assert_eq!(
+            route.sharding_keys().as_slice(),
+            &[
+                Value::Integer(1),
+                Value::String("tenant".to_owned()),
+                Value::Uuid(uuid),
+            ]
+        );
+        assert_eq!(route.shard, i64::from(Shard::Unknown));
+        assert_eq!(route.read_write, u8::from(ReadWrite::Read));
+    }
+
+    #[test]
+    fn empty_sharding_keys_are_valid() {
+        let route = Route::with_sharding_keys(Vec::new(), ReadWrite::Unknown);
+        assert!(route.sharding_keys().is_empty());
     }
 }

pgdog-plugin/src/lib.rs

@@ -35,7 +35,7 @@
 //!
 //! ```toml
 //! [dependencies]
-//! pgdog-plugin = "0.2.0"
+//! pgdog-plugin = "0.5.0"
 //! ```
 //!
 //! # Required methods
@@ -102,6 +102,15 @@
 //! }
 //! ```
 //!
+//! A plugin can also extract typed sharding keys and let PgDog calculate the
+//! destination shards:
+//!
+//! ```
+//! # use pgdog_plugin::prelude::*;
+//! let keys = vec![Value::Integer(42), Value::from("tenant-a")];
+//! let route = Route::with_sharding_keys(keys, ReadWrite::Read);
+//! ```
+//!
 //! ### Parsing parameters
 //!
 //! If your clients are using prepared statements (or the extended protocol), query parameters will be sent separately
@@ -203,13 +212,15 @@ pub mod parameters;
 pub mod plugin;
 pub mod prelude;
 pub mod string;
+pub mod value;
 
 pub use config::Config;
 pub use context::*;
 pub use parameters::*;
 pub use pgdog_postgres_types::Format as ParameterFormat;
 pub use plugin::*;
 pub use string::PdStr;
+pub use value::{ShardingKeys, Value};
 
 pub use libloading;
 

pgdog-plugin/src/prelude.rs

@@ -3,6 +3,6 @@
 #[cfg(feature = "pg_query")]
 pub use crate::pg_query;
 pub use crate::{
-    Context, ParameterFormat, PdStr, Plugin, ReadWrite, Route, Shard,
+    Context, ParameterFormat, PdStr, Plugin, ReadWrite, Route, Shard, ShardingKeys, Value,
     parameters::{Parameter, ParameterValue, Parameters},
 };