feat: add primary only plugin (#826)

```
pgdog=# EXPLAIN SELECT * FROM sessions;
                               QUERY PLAN                               
------------------------------------------------------------------------
 Seq Scan on sessions  (cost=0.00..32.60 rows=2260 width=8)
 
 PgDog Routing:
   Summary: shard=0 role=primary
   Shard 0: SELECT omnishard round robin
   Note: plugin pgdog_primary_only_tables adjusted routing role=primary
(6 rows)
```

#823

Author: levkk

.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.lock

@@ -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"

Cargo.toml

@@ -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

Dockerfile

@@ -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/pgdog.toml

@@ -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/pgx/table_routing_test.go

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

integration/load_balancer/plugin_config.toml

@@ -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"

integration/load_balancer/run.sh

@@ -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-config/src/users.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 {