pgdog.git / summary / log / commit / refs
commit 812ed6bdfcd8a00afecd7591daa2ea4366814e50
Author: Lev Kokotov <lev.kokotov@gmail.com>
Date: Mon Jan 06 03:58:02 2025 +0000
plugins!
Cargo.toml | 7 ++++-
examples/routing-plugin/src/lib.rs | 12 ++++---
pgdog-plugin/build.rs | 2 ++
pgdog-plugin/include/types.h | 14 +++++++++
pgdog-plugin/src/bindings.rs | 25 ++++++++++++++-
pgdog-plugin/src/plugin.rs | 5 +++
pgdog-plugin/src/query.rs | 14 +++++++--
pgdog-plugin/src/route.rs | 14 ++++++++-
pgdog.toml | 4 +++
pgdog/src/backend/pool/connection.rs | 28 ++++++++++++-----
pgdog/src/backend/pool/shard.rs | 8 +++--
pgdog/src/config/mod.rs | 8 +++--
pgdog/src/frontend/buffer.rs | 20 +++++++++++-
pgdog/src/frontend/client.rs | 10 ++++--
pgdog/src/frontend/error.rs | 3 ++
pgdog/src/frontend/listener.rs | 8 +++--
pgdog/src/frontend/mod.rs | 1 +
pgdog/src/frontend/router/error.rs | 13 ++++++++
pgdog/src/frontend/router/mod.rs | 61 ++++++++++++++++++++++++++++++++++++
pgdog/src/main.rs | 2 ++
pgdog/src/net/parameter.rs | 2 +-
pgdog/src/net/stream.rs | 6 ++--
pgdog/src/plugin.rs | 30 ++++++++++++++++--
plugins/pgdog-routing/Cargo.toml | 12 +++++++
plugins/pgdog-routing/src/lib.rs | 44 ++++++++++++++++++++++++++
25 files changed, 321 insertions(+), 32 deletions(-)
diff --git a/Cargo.toml b/Cargo.toml
index 1fb801ef..44f781ee 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,4 +1,9 @@
[workspace]
-members = ["examples/routing-plugin", "pgdog", "pgdog-plugin"]
+members = [
+ "examples/routing-plugin",
+ "pgdog",
+ "pgdog-plugin",
+ "plugins/pgdog-routing",
+]
resolver = "2"
diff --git a/examples/routing-plugin/src/lib.rs b/examples/routing-plugin/src/lib.rs
index a08a10ac..1c38ca9d 100644
--- a/examples/routing-plugin/src/lib.rs
+++ b/examples/routing-plugin/src/lib.rs
@@ -1,17 +1,21 @@
-use pgdog_plugin::{bindings, Query, Route};
+//! Simple routing plugin example using Rust.
+use pgdog_plugin::{
+ bindings::{self, Shard_ANY},
+ Query, Route,
+};
/// Route query.
#[no_mangle]
-pub unsafe extern "C" fn pgdog_route_query(query: bindings::Query) -> Route {
+pub extern "C" fn pgdog_route_query(query: bindings::Query) -> Route {
let query = Query::from(query);
if query.query().to_lowercase().starts_with("select") {
Route {
- shard: -1, // Any shard.
+ shard: Shard_ANY, // Any shard.
affinity: bindings::Affinity_READ,
}
} else {
Route {
- shard: -1, // Any shard.
+ shard: Shard_ANY, // Any shard.
affinity: bindings::Affinity_WRITE,
}
}
diff --git a/pgdog-plugin/build.rs b/pgdog-plugin/build.rs
index d78d5a5c..9c319567 100644
--- a/pgdog-plugin/build.rs
+++ b/pgdog-plugin/build.rs
@@ -1,6 +1,8 @@
use std::path::PathBuf;
fn main() {
+ println!("cargo:rerun-if-changed=include/types.h");
+
let bindings = bindgen::Builder::default()
.header("include/wrapper.h")
// Tell cargo to invalidate the built crate whenever any of the
diff --git a/pgdog-plugin/include/types.h b/pgdog-plugin/include/types.h
index 219d5595..c52c2f17 100644
--- a/pgdog-plugin/include/types.h
+++ b/pgdog-plugin/include/types.h
@@ -1,16 +1,30 @@
/*
* Query.
*/
+typedef struct Value {
+ int len;
+ const char *data;
+ int oid;
+} Value;
+
typedef struct Query {
int len;
const char *query;
+ int num_values;
+ const Value *values;
} Query;
typedef enum Affinity {
READ = 1,
WRITE = 2,
+ UNKNOWN = 3,
} Affinity;
+typedef enum Shard {
+ ANY = -1,
+ ALL = -2,
+} Shard;
+
typedef struct Route {
Affinity affinity;
int shard;
diff --git a/pgdog-plugin/src/bindings.rs b/pgdog-plugin/src/bindings.rs
index b42de011..04334b8f 100644
--- a/pgdog-plugin/src/bindings.rs
+++ b/pgdog-plugin/src/bindings.rs
@@ -1,21 +1,44 @@
/* automatically generated by rust-bindgen 0.71.1 */
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct Value {
+ pub len: ::std::os::raw::c_int,
+ pub data: *const ::std::os::raw::c_char,
+ pub oid: ::std::os::raw::c_int,
+}
+#[allow(clippy::unnecessary_operation, clippy::identity_op)]
+const _: () = {
+ ["Size of Value"][::std::mem::size_of::<Value>() - 24usize];
+ ["Alignment of Value"][::std::mem::align_of::<Value>() - 8usize];
+ ["Offset of field: Value::len"][::std::mem::offset_of!(Value, len) - 0usize];
+ ["Offset of field: Value::data"][::std::mem::offset_of!(Value, data) - 8usize];
+ ["Offset of field: Value::oid"][::std::mem::offset_of!(Value, oid) - 16usize];
+};
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct Query {
pub len: ::std::os::raw::c_int,
pub query: *const ::std::os::raw::c_char,
+ pub num_values: ::std::os::raw::c_int,
+ pub values: *const Value,
}
#[allow(clippy::unnecessary_operation, clippy::identity_op)]
const _: () = {
- ["Size of Query"][::std::mem::size_of::<Query>() - 16usize];
+ ["Size of Query"][::std::mem::size_of::<Query>() - 32usize];
["Alignment of Query"][::std::mem::align_of::<Query>() - 8usize];
["Offset of field: Query::len"][::std::mem::offset_of!(Query, len) - 0usize];
["Offset of field: Query::query"][::std::mem::offset_of!(Query, query) - 8usize];
+ ["Offset of field: Query::num_values"][::std::mem::offset_of!(Query, num_values) - 16usize];
+ ["Offset of field: Query::values"][::std::mem::offset_of!(Query, values) - 24usize];
};
pub const Affinity_READ: Affinity = 1;
pub const Affinity_WRITE: Affinity = 2;
+pub const Affinity_UNKNOWN: Affinity = 3;
pub type Affinity = ::std::os::raw::c_uint;
+pub const Shard_ANY: Shard = -1;
+pub const Shard_ALL: Shard = -2;
+pub type Shard = ::std::os::raw::c_int;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct Route {
diff --git a/pgdog-plugin/src/plugin.rs b/pgdog-plugin/src/plugin.rs
index 6385d729..d13888b2 100644
--- a/pgdog-plugin/src/plugin.rs
+++ b/pgdog-plugin/src/plugin.rs
@@ -43,4 +43,9 @@ impl<'a> Plugin<'a> {
pub fn name(&self) -> &str {
&self.name
}
+
+ /// Check that we have the required methods.
+ pub fn valid(&self) -> bool {
+ self.route.is_some()
+ }
}
diff --git a/pgdog-plugin/src/query.rs b/pgdog-plugin/src/query.rs
index a839d341..b2bc0100 100644
--- a/pgdog-plugin/src/query.rs
+++ b/pgdog-plugin/src/query.rs
@@ -1,4 +1,4 @@
-use crate::bindings;
+use crate::bindings::{self};
use std::ffi::{c_char, CStr, CString, NulError};
use std::marker::PhantomData;
@@ -12,7 +12,11 @@ pub struct Query<'a> {
len: usize,
/// Query string.
query: *const c_char,
- /// Lifetime marker.
+ /// Number of parameters if any.
+ num_values: usize,
+ values: *const bindings::Value,
+ /// Lifetime marker ensuring that the CString
+ /// from which this query is created is not deallocated too soon.
_lifetime: PhantomData<&'a ()>,
}
@@ -27,6 +31,8 @@ impl From<Query<'_>> for bindings::Query {
Self {
len: value.len as i32,
query: value.query as *mut i8,
+ num_values: 0,
+ values: null(),
}
}
}
@@ -36,6 +42,8 @@ impl From<bindings::Query> for Query<'_> {
Self {
len: value.len as usize,
query: value.query as *const c_char,
+ num_values: 0,
+ values: null(),
_lifetime: PhantomData,
}
}
@@ -53,6 +61,8 @@ impl<'a> Query<'a> {
Self {
len: query.as_bytes().len(),
query: query.as_ptr() as *const c_char,
+ num_values: 0,
+ values: null(),
_lifetime: PhantomData,
}
}
diff --git a/pgdog-plugin/src/route.rs b/pgdog-plugin/src/route.rs
index f71a105e..56b6e54b 100644
--- a/pgdog-plugin/src/route.rs
+++ b/pgdog-plugin/src/route.rs
@@ -1,6 +1,14 @@
use crate::bindings::*;
impl Route {
+ ///
+ pub fn unknown() -> Route {
+ Route {
+ shard: Shard_ANY,
+ affinity: Affinity_UNKNOWN,
+ }
+ }
+
/// Is this a read?
pub fn read(&self) -> bool {
self.affinity == Affinity_READ
@@ -20,8 +28,12 @@ impl Route {
}
}
+ pub fn any_shard(&self) -> bool {
+ self.shard == Shard_ANY
+ }
+
/// Query should go across all shards.
pub fn cross_shard(&self) -> bool {
- self.shard == -2
+ self.shard == Shard_ALL
}
}
diff --git a/pgdog.toml b/pgdog.toml
new file mode 100644
index 00000000..4ca08477
--- /dev/null
+++ b/pgdog.toml
@@ -0,0 +1,4 @@
+
+[general]
+
+plugins = ["pgdog_routing"]
diff --git a/pgdog/src/backend/pool/connection.rs b/pgdog/src/backend/pool/connection.rs
index b6b7eebc..6b305d22 100644
--- a/pgdog/src/backend/pool/connection.rs
+++ b/pgdog/src/backend/pool/connection.rs
@@ -1,5 +1,6 @@
//! Server connection.
+use pgdog_plugin::Route;
use tokio::time::sleep;
use crate::{
@@ -10,8 +11,9 @@ use crate::{
use super::{
super::{pool::Guard, Error},
- Cluster,
+ Address, Cluster,
};
+
use std::time::Duration;
/// Wrapper around a server connection.
@@ -48,13 +50,13 @@ impl Connection {
}
/// Create a server connection if one doesn't exist already.
- pub async fn connect(&mut self, id: &BackendKeyData) -> Result<(), Error> {
+ pub async fn connect(&mut self, id: &BackendKeyData, route: &Route) -> Result<(), Error> {
if self.server.is_none() && self.admin.is_none() {
- match self.try_conn(id).await {
+ match self.try_conn(id, route).await {
Ok(()) => (),
Err(Error::Pool(super::Error::Offline)) => {
self.reload()?;
- return Ok(self.try_conn(id).await?);
+ return Ok(self.try_conn(id, route).await?);
}
Err(err) => return Err(err.into()),
}
@@ -63,8 +65,15 @@ impl Connection {
Ok(())
}
- async fn try_conn(&mut self, id: &BackendKeyData) -> Result<(), Error> {
- let server = self.cluster()?.primary(0, id).await?;
+ async fn try_conn(&mut self, id: &BackendKeyData, route: &Route) -> Result<(), Error> {
+ let shard = route.shard().unwrap_or(0);
+
+ let server = if route.read() {
+ self.cluster()?.replica(shard, id).await?
+ } else {
+ self.cluster()?.primary(shard, id).await?
+ };
+
self.server = Some(server);
Ok(())
@@ -75,7 +84,7 @@ impl Connection {
if self.admin.is_some() {
Ok(ParameterStatus::fake())
} else {
- self.connect(id).await?;
+ self.connect(id, &Route::unknown()).await?;
let params = self
.server()?
.params()
@@ -133,6 +142,11 @@ impl Connection {
}
}
+ /// Get connected server address.
+ pub fn addr(&mut self) -> Result<&Address, Error> {
+ Ok(self.server()?.addr())
+ }
+
#[inline]
fn cluster(&self) -> Result<&Cluster, Error> {
Ok(self.cluster.as_ref().ok_or(Error::NotConnected)?)
diff --git a/pgdog/src/backend/pool/shard.rs b/pgdog/src/backend/pool/shard.rs
index 53bc9661..e14b56b6 100644
--- a/pgdog/src/backend/pool/shard.rs
+++ b/pgdog/src/backend/pool/shard.rs
@@ -25,9 +25,13 @@ impl Shard {
self.primary.get(id).await
}
- /// Get a connection to a shard replica.
+ /// Get a connection to a shard replica, if any.
pub async fn replica(&self, id: &BackendKeyData) -> Result<Guard, Error> {
- self.replicas.get(id).await
+ if self.replicas.is_empty() {
+ self.primary.get(id).await
+ } else {
+ self.replicas.get(id).await
+ }
}
/// Create new identical connection pool.
diff --git a/pgdog/src/config/mod.rs b/pgdog/src/config/mod.rs
index 459dcde1..fe10c6b2 100644
--- a/pgdog/src/config/mod.rs
+++ b/pgdog/src/config/mod.rs
@@ -22,8 +22,10 @@ pub fn config() -> Arc<Config> {
/// Load the configuration file from disk.
pub fn load() -> Result<Config, Error> {
if let Ok(config) = read_to_string("pgdog.toml") {
- info!("Loading pgdog.toml");
- Ok(toml::from_str(&config)?)
+ let config: Config = toml::from_str(&config)?;
+ CONFIG.store(Arc::new(config.clone()));
+ info!("Loaded pgdog.toml");
+ Ok(config)
} else {
info!("Loading default configuration");
Ok(Config::default())
@@ -52,6 +54,8 @@ pub struct General {
pub port: u16,
#[serde(default = "General::workers")]
pub workers: usize,
+ #[serde(default)]
+ pub plugins: Vec<String>,
}
impl General {
diff --git a/pgdog/src/frontend/buffer.rs b/pgdog/src/frontend/buffer.rs
index 3ba61657..49f0d40a 100644
--- a/pgdog/src/frontend/buffer.rs
+++ b/pgdog/src/frontend/buffer.rs
@@ -2,7 +2,10 @@
use std::ops::{Deref, DerefMut};
-use crate::net::messages::{Message, Protocol};
+use crate::net::{
+ messages::{parse::Parse, FromBytes, Message, Protocol, Query, ToBytes},
+ Error,
+};
/// Message buffer.
pub struct Buffer {
@@ -57,6 +60,21 @@ impl Buffer {
pub fn len(&self) -> usize {
self.buffer.iter().map(|b| b.len()).sum()
}
+
+ /// If this buffer contains a query, retrive it.
+ pub fn query(&self) -> Result<Option<String>, Error> {
+ for message in &self.buffer {
+ if message.code() == 'Q' {
+ let query = Query::from_bytes(message.to_bytes()?)?;
+ return Ok(Some(query.query));
+ } else if message.code() == 'P' {
+ let parse = Parse::from_bytes(message.to_bytes()?)?;
+ return Ok(Some(parse.query));
+ }
+ }
+
+ Ok(None)
+ }
}
impl From<Buffer> for Vec<Message> {
diff --git a/pgdog/src/frontend/client.rs b/pgdog/src/frontend/client.rs
index b35f4247..09d82aed 100644
--- a/pgdog/src/frontend/client.rs
+++ b/pgdog/src/frontend/client.rs
@@ -1,8 +1,9 @@
//! Frontend client.
use tokio::select;
+use tracing::debug;
-use super::{Buffer, Error};
+use super::{Buffer, Error, Router};
use crate::backend::pool::Connection;
use crate::net::messages::{Authentication, BackendKeyData, Protocol, ReadyForQuery};
use crate::net::{parameter::Parameters, Stream};
@@ -63,6 +64,7 @@ impl Client {
let admin = database == "admin";
let mut backend = Connection::new(user, database, admin)?;
+ let mut router = Router::new();
let mut flush = false;
self.state = State::Idle;
@@ -77,9 +79,13 @@ impl Client {
flush = buffer.flush();
if !backend.connected() {
+ router.query(&buffer)?;
+
self.state = State::Waiting;
- backend.connect(&self.id).await?;
+ backend.connect(&self.id, router.route()).await?;
self.state = State::Active;
+
+ debug!("client paired with {}", backend.addr()?);
}
backend.send(buffer.into()).await?;
diff --git a/pgdog/src/frontend/error.rs b/pgdog/src/frontend/error.rs
index d33d0e22..16dd202b 100644
--- a/pgdog/src/frontend/error.rs
+++ b/pgdog/src/frontend/error.rs
@@ -16,4 +16,7 @@ pub enum Error {
#[error("\"{0}\" parameter is missing")]
Parameter(String),
+
+ #[error("{0}")]
+ Router(#[from] super::router::Error),
}
diff --git a/pgdog/src/frontend/listener.rs b/pgdog/src/frontend/listener.rs
index ab9e4d5f..9fa23fe9 100644
--- a/pgdog/src/frontend/listener.rs
+++ b/pgdog/src/frontend/listener.rs
@@ -54,8 +54,12 @@ impl Listener {
let clients = self.clients.clone();
tokio::spawn(async move {
- Self::handle_client(stream, addr, clients).await?;
- Ok::<(), Error>(())
+ match Self::handle_client(stream, addr, clients).await {
+ Ok(_) => (),
+ Err(err) => {
+ error!("client crashed: {:?}", err);
+ }
+ }
});
}
diff --git a/pgdog/src/frontend/mod.rs b/pgdog/src/frontend/mod.rs
index 025edca4..8348515d 100644
--- a/pgdog/src/frontend/mod.rs
+++ b/pgdog/src/frontend/mod.rs
@@ -9,3 +9,4 @@ pub mod router;
pub use buffer::Buffer;
pub use client::Client;
pub use error::Error;
+pub use router::Router;
diff --git a/pgdog/src/frontend/router/error.rs b/pgdog/src/frontend/router/error.rs
new file mode 100644
index 00000000..d4dd6dc0
--- /dev/null
+++ b/pgdog/src/frontend/router/error.rs
@@ -0,0 +1,13 @@
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum Error {
+ #[error("routing plugin missing")]
+ RoutingPluginMissing,
+
+ #[error("plugin error")]
+ PluginError(#[from] std::ffi::NulError),
+
+ #[error("no query in buffer")]
+ NoQueryInBuffer,
+}
diff --git a/pgdog/src/frontend/router/mod.rs b/pgdog/src/frontend/router/mod.rs
index 80a98eb2..df001421 100644
--- a/pgdog/src/frontend/router/mod.rs
+++ b/pgdog/src/frontend/router/mod.rs
@@ -1,3 +1,64 @@
//! Query router.
+use crate::plugin::plugins;
+
+use pgdog_plugin::{FfiQuery, Route};
+use tokio::time::Instant;
+use tracing::debug;
+
+pub mod error;
pub mod parser;
+
+pub use error::Error;
+
+use super::Buffer;
+
+/// Query router.
+pub struct Router {
+ route: Route,
+}
+
+impl Router {
+ /// Create new router.
+ pub fn new() -> Router {
+ Self {
+ route: Route::unknown(),
+ }
+ }
+
+ /// Route a query to a shard.
+ pub fn query(&mut self, buffer: &Buffer) -> Result<Route, Error> {
+ let query = buffer
+ .query()
+ .map_err(|_| Error::NoQueryInBuffer)?
+ .ok_or(Error::NoQueryInBuffer)?;
+ let query = FfiQuery::new(&query)?;
+ let now = Instant::now();
+
+ for plugin in plugins() {
+ match plugin.route(query.query()) {
+ None => continue,
+ Some(route) => {
+ self.route = route;
+
+ debug!(
+ "routing {} to shard {} [{}, {:.3}ms]",
+ if route.read() { "read" } else { "write" },
+ route.shard().unwrap_or(0),
+ plugin.name(),
+ now.elapsed().as_secs_f64() * 1000.0,
+ );
+
+ return Ok(route);
+ }
+ }
+ }
+
+ Ok(Route::unknown())
+ }
+
+ /// Get current route.
+ pub fn route(&self) -> &Route {
+ &self.route
+ }
+}
diff --git a/pgdog/src/main.rs b/pgdog/src/main.rs
index 946a751e..5afb7ab2 100644
--- a/pgdog/src/main.rs
+++ b/pgdog/src/main.rs
@@ -28,6 +28,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = load()?;
+ plugin::load_from_config()?;
+
let runtime = match config.general.workers {
0 => {
let mut binding = Builder::new_current_thread();
diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs
index dd7e49eb..27768b14 100644
--- a/pgdog/src/net/parameter.rs
+++ b/pgdog/src/net/parameter.rs
@@ -35,7 +35,7 @@ impl Parameters {
}
/// Get parameter value or returned a default value if it doesn't exist.
- pub fn get_default<'a>(&'a self, name: &str, default_value: &'a str) -> &str {
+ pub fn get_default<'a>(&'a self, name: &str, default_value: &'a str) -> &'a str {
self.get(name).map_or(default_value, |p| p)
}
}
diff --git a/pgdog/src/net/stream.rs b/pgdog/src/net/stream.rs
index f9183b44..65526d55 100644
--- a/pgdog/src/net/stream.rs
+++ b/pgdog/src/net/stream.rs
@@ -4,7 +4,7 @@ use bytes::{BufMut, BytesMut};
use pin_project::pin_project;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream, ReadBuf};
use tokio::net::TcpStream;
-use tracing::debug;
+use tracing::{debug, trace};
use std::io::Error;
use std::pin::Pin;
@@ -87,7 +87,7 @@ impl Stream {
/// This is fast because the stream is buffered. Make sure to call [`Stream::send_flush`]
/// for the last message in the exchange.
pub async fn send(&mut self, message: impl Protocol) -> Result<usize, crate::net::Error> {
- debug!("📡 <= {}", message.code());
+ trace!("📡 <= {}", message.code());
let bytes = message.to_bytes()?;
match self {
@@ -152,7 +152,7 @@ impl Stream {
let message = Message::new(bytes.freeze());
- debug!("📡 => {}", message.code());
+ trace!("📡 => {}", message.code());
Ok(message)
}
diff --git a/pgdog/src/plugin.rs b/pgdog/src/plugin.rs
index ebab5ee3..0d31126d 100644
--- a/pgdog/src/plugin.rs
+++ b/pgdog/src/plugin.rs
@@ -4,7 +4,7 @@ use once_cell::sync::OnceCell;
use pgdog_plugin::libloading;
use pgdog_plugin::libloading::Library;
use pgdog_plugin::Plugin;
-use tracing::{error, info};
+use tracing::{error, info, warn};
static LIBS: OnceCell<Vec<Library>> = OnceCell::new();
pub static PLUGINS: OnceCell<Vec<Plugin>> = OnceCell::new();
@@ -35,8 +35,13 @@ pub fn load(names: &[&str]) -> Result<(), libloading::Error> {
for (i, name) in names.iter().enumerate() {
if let Some(lib) = LIBS.get().unwrap().get(i) {
let plugin = Plugin::load(name, lib);
- plugins.push(plugin);
- info!("Loaded \"{}\" plugin", name);
+
+ if !plugin.valid() {
+ warn!("plugin \"{}\" is missing required symbols, skipping", name);
+ } else {
+ plugins.push(plugin);
+ info!("Loaded \"{}\" plugin", name);
+ }
}
}
@@ -56,6 +61,25 @@ pub fn plugin(name: &str) -> Option<&Plugin> {
None
}
+/// Get all loaded plugins.
+pub fn plugins() -> &'static Vec<Plugin<'static>> {
+ PLUGINS.get().unwrap()
+}
+
+/// Load plugins from config.
+pub fn load_from_config() -> Result<(), libloading::Error> {
+ let config = crate::config::config();
+
+ let plugins = &config
+ .general
+ .plugins
+ .iter()
+ .map(|s| s.as_str())
+ .collect::<Vec<_>>();
+
+ load(plugins)
+}
+
#[cfg(test)]
mod test {
use pgdog_plugin::FfiQuery;
diff --git a/plugins/pgdog-routing/Cargo.toml b/plugins/pgdog-routing/Cargo.toml
new file mode 100644
index 00000000..5fc4f4a1
--- /dev/null
+++ b/plugins/pgdog-routing/Cargo.toml
@@ -0,0 +1,12 @@
+[package]
+name = "pgdog-routing"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+pgdog-plugin = { path = "../../pgdog-plugin" }
+pg_query = "6.0"
+tracing = "*"
+
+[lib]
+crate-type = ["rlib", "cdylib"]
diff --git a/plugins/pgdog-routing/src/lib.rs b/plugins/pgdog-routing/src/lib.rs
new file mode 100644
index 00000000..4ae6fd6e
--- /dev/null
+++ b/plugins/pgdog-routing/src/lib.rs
@@ -0,0 +1,44 @@
+//! Parse queries using pg_query and route all SELECT queries
+//! to replicas. All other queries are routed to a primary.
+use pg_query::{parse, NodeEnum};
+use pgdog_plugin::bindings::Shard_ANY;
+use pgdog_plugin::{bindings, Affinity_READ, Affinity_WRITE};
+use pgdog_plugin::{Query, Route};
+
+#[no_mangle]
+pub extern "C" fn pgdog_route_query(query: bindings::Query) -> Route {
+ let query = Query::from(query);
+ match route_internal(query.query()) {
+ Ok(route) => route,
+ Err(_) => Route::unknown(),
+ }
+}
+
+fn route_internal(query: &str) -> Result<Route, pg_query::Error> {
+ let ast = parse(query)?;
+
+ match ast.protobuf.stmts.first() {
+ Some(query) => match query.stmt {
+ Some(ref node) => match node.node {
+ Some(NodeEnum::SelectStmt(ref _stmt)) => {
+ return Ok(Route {
+ affinity: Affinity_READ,
+ shard: Shard_ANY,
+ });
+ }
+
+ Some(_) => (),
+
+ None => (),
+ },
+ None => (),
+ },
+
+ None => (),
+ }
+
+ Ok(Route {
+ affinity: Affinity_WRITE,
+ shard: Shard_ANY,
+ })
+}
[parent: f6356e514cc4]