feat: add primary only plugin

.schema/pgdog.schema.json

@@ -1256,6 +1256,13 @@
       "description": "Plugins are dynamically loaded at PgDog startup. These settings control which plugins are loaded.\n\nNote: Plugins can only be configured at PgDog startup. They cannot be changed after the process is running.\n\nhttps://docs.pgdog.dev/configuration/pgdog.toml/plugins/",
       "type": "object",
       "properties": {
+        "config": {
+          "description": "Path to the configuration file for the plugin, if any. Plugin-specific settings can be\nplaced there. It's completely plugin-specific and any fomrat is acceptable.",
+          "type": [
+            "string",
+            "null"
+          ]
+        },
         "name": {
           "description": "Name of the plugin to load. This is used by PgDog to look up the shared library object in `LD_LIBRARY_PATH`. For example, if your plugin name is `router`, PgDog will look for `librouter.so` on Linux, `librouter.dll` on Windows, and `librouter.dylib` on Mac OS.\n\n**Note:** Make sure the user running PgDog has read & execute permissions on the library.\n\nhttps://docs.pgdog.dev/configuration/pgdog.toml/plugins/#name",
           "type": "string"

Cargo.toml

@@ -10,15 +10,15 @@ members = [
   "pgdog-postgres-types",
   "pgdog-stats",
   "pgdog-vector",
-  "plugins/pgdog-example-plugin",
+  "plugins/pgdog-example-plugin", "plugins/pgdog-primary-only-tables",
   "scripts/*",
 ]
 
 [workspace.package]
 edition = "2024"
 
 [workspace.dependencies]
-pgdog-plugin = { path = "./pgdog-plugin", version = "0.2.0" }
+pgdog-plugin = { path = "./pgdog-plugin", version = "0.3.0" }
 pgdog-config = { path = "./pgdog-config", version = "0.1.0" }
 schemars = { version = "1.2.1", features = ["uuid1"] }
 serde_json = "1.0"

Dockerfile

@@ -16,7 +16,9 @@ RUN source ~/.cargo/env && \
         export RUSTFLAGS="-Ctarget-feature=+lse"; \
     fi && \
     cd pgdog && \
-    cargo build --release
+    cargo build --release && \
+    cd .. && \
+    cargo build --release -p pgdog-primary-only-tables
 
 FROM ubuntu:latest
 ENV RUST_LOG=info
@@ -34,6 +36,7 @@ RUN install -d /usr/share/postgresql-common/pgdg && \
 RUN apt update && apt install -y postgresql-client-${PSQL_VERSION}
 
 COPY --from=builder /build/target/release/pgdog /usr/local/bin/pgdog
+COPY --from=builder /build/target/release/libpgdog_primary_only_tables.so /usr/lib/libpgdog_primary_only_tables.so
 
 WORKDIR /pgdog
 STOPSIGNAL SIGINT

integration/load_balancer/pgdog.toml

@@ -16,6 +16,7 @@ pooler_mode = "transaction"
 load_balancing_strategy = "round_robin"
 auth_type = "trust"
 read_write_split = "exclude_primary"
+expanded_explain = true
 lsn_check_delay = 0
 
 [rewrite]
@@ -50,6 +51,10 @@ role = "replica"
 port = 45002
 
 
+[[plugins]]
+name = "pgdog_primary_only_tables"
+config = "integration/load_balancer/plugin_config.toml"
+
 [tcp]
 retries = 3
 time = 1000

integration/load_balancer/pgx/table_routing_test.go

@@ -0,0 +1,68 @@
+package main
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPrimaryOnlyTables(t *testing.T) {
+	pool := GetPool()
+	defer pool.Close()
+
+	_, err := pool.Exec(context.Background(), `CREATE TABLE IF NOT EXISTS lb_plugin_primary_only (
+		id BIGINT,
+		data VARCHAR
+	)`)
+	assert.NoError(t, err)
+	defer pool.Exec(context.Background(), "DROP TABLE IF EXISTS lb_plugin_primary_only")
+
+	time.Sleep(2 * time.Second)
+
+	ResetStats()
+
+	for i := range 50 {
+		_, err = pool.Exec(context.Background(), "SELECT $1::bigint FROM lb_plugin_primary_only", int64(i))
+		assert.NoError(t, err)
+	}
+
+	primaryCalls := LoadStatsForPrimary("lb_plugin_primary_only")
+	assert.Equal(t, int64(50), primaryCalls.Calls)
+}
+
+func TestPrimaryOnlyTablesExplain(t *testing.T) {
+	pool := GetPool()
+	defer pool.Close()
+
+	_, err := pool.Exec(context.Background(), `CREATE TABLE IF NOT EXISTS lb_plugin_primary_only (
+		id BIGINT,
+		data VARCHAR
+	)`)
+	assert.NoError(t, err)
+	defer pool.Exec(context.Background(), "DROP TABLE IF EXISTS lb_plugin_primary_only")
+
+	time.Sleep(2 * time.Second)
+
+	rows, err := pool.Query(context.Background(), "EXPLAIN SELECT * FROM lb_plugin_primary_only")
+	assert.NoError(t, err)
+
+	var explainLines []string
+	foundPluginAnnotation := false
+	for rows.Next() {
+		var line string
+		err = rows.Scan(&line)
+		assert.NoError(t, err)
+		explainLines = append(explainLines, line)
+
+		if strings.Contains(line, "plugin pgdog_primary_only_tables adjusted routing role=primary") {
+			foundPluginAnnotation = true
+		}
+	}
+	rows.Close()
+
+	t.Logf("EXPLAIN output:\n%s", strings.Join(explainLines, "\n"))
+	assert.True(t, foundPluginAnnotation, "EXPLAIN output should contain plugin routing annotation")
+}

integration/load_balancer/plugin_config.toml

@@ -0,0 +1,2 @@
+[[tables]]
+name = "lb_plugin_primary_only"

integration/load_balancer/run.sh

@@ -27,6 +27,13 @@ for p in 45000 45001 45002; do
     fi
 done
 
+pushd ${SCRIPT_DIR}/../../plugins/pgdog-primary-only-tables
+cargo build --release
+popd
+
+export LD_LIBRARY_PATH=${SCRIPT_DIR}/../../target/release:${LD_LIBRARY_PATH:-}
+export DYLD_LIBRARY_PATH=${LD_LIBRARY_PATH}
+
 docker-compose up -d
 
 echo "Waiting for Postgres to be ready"

pgdog-config/src/users.rs

@@ -1,5 +1,6 @@
 use serde::{Deserialize, Serialize};
 use std::env;
+use std::path::PathBuf;
 use tracing::warn;
 
 use super::core::Config;
@@ -21,6 +22,10 @@ pub struct Plugin {
     ///
     /// https://docs.pgdog.dev/configuration/pgdog.toml/plugins/#name
     pub name: String,
+
+    /// Path to the configuration file for the plugin, if any. Plugin-specific settings can be
+    /// placed there. It's completely plugin-specific and any fomrat is acceptable.
+    pub config: Option<PathBuf>,
 }
 
 /// This configuration controls which users are allowed to connect to PgDog. This is a TOML list so for each user, add a `[[users]]` section to `users.toml`.

pgdog-macros/src/lib.rs

@@ -40,6 +40,11 @@ pub fn plugin(_input: TokenStream) -> TokenStream {
                 *output = version;
             }
         }
+
+        #[unsafe(no_mangle)]
+        pub extern "C" fn pgdog_logging_init(config: pgdog_plugin::PdConfig) {
+            pgdog_plugin::logging::init(&config);
+        }
     };
     TokenStream::from(expanded)
 }
@@ -63,6 +68,28 @@ pub fn init(_attr: TokenStream, item: TokenStream) -> TokenStream {
     TokenStream::from(expanded)
 }
 
+/// Generate the `pgdog_config` method that's executed at plugin load time.
+#[proc_macro_attribute]
+pub fn config(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input_fn = parse_macro_input!(item as ItemFn);
+    let fn_name = &input_fn.sig.ident;
+
+    let expanded = quote! {
+
+        #[unsafe(no_mangle)]
+        pub extern "C" fn pgdog_config(
+            pd_config: pgdog_plugin::PdConfig,
+            result: *mut u8)
+        {
+            #input_fn
+
+            #fn_name(pd_config, result);
+        }
+    };
+
+    TokenStream::from(expanded)
+}
+
 /// Generate the `pgdog_fini` method that runs at PgDog shutdown.
 #[proc_macro_attribute]
 pub fn fini(_attr: TokenStream, item: TokenStream) -> TokenStream {

pgdog-plugin/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pgdog-plugin"
-version = "0.2.0"
+version = "0.3.0"
 edition = "2021"
 license = "MIT"
 authors = ["Lev Kokotov <lev.kokotov@gmail.com>"]
@@ -18,5 +18,8 @@ libloading = "0.8"
 pg_query = { git = "https://github.com/pgdogdev/pg_query.rs.git", rev = "f8c216209f90525f065b47ffde9eb5da803d2dc6" }
 pgdog-macros = { path = "../pgdog-macros", version = "0.1.1" }
 
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
+
 [build-dependencies]
 bindgen = "0.71.0"

pgdog-plugin/include/types.h

@@ -64,6 +64,18 @@ typedef struct PdRouterContext {
 /**
  * Routing decision returned by the plugin.
  */
+/**
+ * Configuration passed from PgDog to plugins.
+ */
+typedef struct PdConfig {
+    /** Log level, e.g. "info", "debug", "error". */
+    struct PdStr log_level;
+    /** Whether to use JSON log format. `1` = `true`, `0` = `false`. */
+    uint8_t log_json;
+    /** Path to the plugin's configuration file. Empty string if not set. */
+    struct PdStr plugin_config;
+} PdConfig;
+
  typedef struct PdRoute {
      /** Which shard the query should go to.
       *

pgdog-plugin/src/lib.rs

@@ -169,6 +169,7 @@ pub mod bindings {
 pub mod ast;
 pub mod comp;
 pub mod context;
+pub mod logging;
 pub mod parameters;
 pub mod plugin;
 pub mod prelude;

pgdog-plugin/src/logging.rs

@@ -0,0 +1,47 @@
+//! Logger initialization for plugins.
+//!
+//! Called automatically by the `#[config]` macro before the plugin's
+//! config function runs, so plugins share PgDog's log level and format.
+
+use std::io::IsTerminal;
+
+use tracing::level_filters::LevelFilter;
+use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
+
+use crate::PdConfig;
+
+pub fn init(config: &PdConfig) {
+    let log_level = std::ops::Deref::deref(&config.log_level);
+    let log_json = config.log_json == 1;
+
+    let filter = EnvFilter::builder()
+        .with_default_directive(log_level.parse().unwrap_or(LevelFilter::INFO.into()))
+        .from_env_lossy();
+
+    if log_json {
+        let format = fmt::layer()
+            .json()
+            .with_ansi(false)
+            .with_writer(std::io::stderr)
+            .with_file(false);
+        #[cfg(not(debug_assertions))]
+        let format = format.with_target(false);
+
+        let _ = tracing_subscriber::registry()
+            .with(format)
+            .with(filter)
+            .try_init();
+    } else {
+        let format = fmt::layer()
+            .with_ansi(std::io::stderr().is_terminal())
+            .with_writer(std::io::stderr)
+            .with_file(false);
+        #[cfg(not(debug_assertions))]
+        let format = format.with_target(false);
+
+        let _ = tracing_subscriber::registry()
+            .with(format)
+            .with(filter)
+            .try_init();
+    }
+}