From b57dbd0fa0c6a5f6d299be21e4985c99186625f1 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 8 Mar 2026 22:12:01 +0000 Subject: [PATCH] feat: add pg_replication module for reactive WAL-based cache invalidation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements a PostgreSQL logical replication module that streams WAL changes, buffers them per XID for transaction atomicity, and publishes CacheKey invalidation events to the existing framework broadcast channel. Key design choices: - Transaction atomicity: changes are buffered until COMMIT so the fanout always receives a consistent view. A 1000-row bulk-insert emits one table-level CacheKey, not 1000 row events. - Native protocol: uses a self-contained raw TCP + postgres-protocol implementation for the replication connection (tokio-postgres 0.7 does not expose copy_both_simple publicly). - Corrected pgoutput v1 parser: original maybe_sql_integration code incorrectly read XID from DML messages; DML messages carry no XID in proto v1 — only Begin does. - Clean integration: publishes CacheKey::Channel events to store.events so the existing fanout task picks them up with zero changes to the fanout loop. - New cx.subscribe(key) API on RenderContext for manually registering subscription keys (needed when queries go through PgPool, not cx.run). New files: src/pg_replication/mod.rs — PgReplicationListener, key helpers src/pg_replication/wal_parser.rs — pgoutput v1 binary protocol parser src/pg_replication/emitter.rs — WalEmitter (tx-buffering + CacheKey emit) src/pg_replication/proto.rs — raw TCP PG wire-protocol connection examples/pg_replication/main.rs — live message board example docker-compose.yml — PG 16 container with wal_level=logical Run the example: docker compose up -d cargo run --example pg_replication --features pg_replication https://claude.ai/code/session_01SLGgXeqKV2o7KTmCaQZZPg --- Cargo.lock | 335 +++++++++++++++++++- Cargo.toml | 22 +- docker-compose.yml | 30 ++ examples/pg_replication/main.rs | 387 +++++++++++++++++++++++ src/context.rs | 24 ++ src/lib.rs | 9 + src/pg_replication/emitter.rs | 311 +++++++++++++++++++ src/pg_replication/mod.rs | 265 ++++++++++++++++ src/pg_replication/proto.rs | 341 +++++++++++++++++++++ src/pg_replication/wal_parser.rs | 505 +++++++++++++++++++++++++++++++ 10 files changed, 2212 insertions(+), 17 deletions(-) create mode 100644 docker-compose.yml create mode 100644 examples/pg_replication/main.rs create mode 100644 src/pg_replication/emitter.rs create mode 100644 src/pg_replication/mod.rs create mode 100644 src/pg_replication/proto.rs create mode 100644 src/pg_replication/wal_parser.rs diff --git a/Cargo.lock b/Cargo.lock index cfaa446..0c9a6c3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -370,6 +379,12 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -519,7 +534,19 @@ checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" dependencies = [ "cfg-if", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", ] [[package]] @@ -530,7 +557,7 @@ checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" dependencies = [ "cfg-if", "libc", - "r-efi", + "r-efi 6.0.0", "wasip2", "wasip3", ] @@ -909,16 +936,20 @@ version = "0.1.0" dependencies = [ "anyhow", "axum", + "bytes", "chrono", "dashmap", "futures", "maud", + "postgres-protocol", "serde", "serde_json", "sqlx", "tokio", + "tokio-postgres", "tokio-stream", "tracing", + "tracing-subscriber", "uuid", ] @@ -937,6 +968,15 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + [[package]] name = "matchit" version = "0.7.3" @@ -996,7 +1036,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" dependencies = [ "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", + "windows-sys 0.61.2", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ "windows-sys 0.61.2", ] @@ -1011,7 +1060,7 @@ dependencies = [ "num-integer", "num-iter", "num-traits", - "rand", + "rand 0.8.5", "smallvec", "zeroize", ] @@ -1046,6 +1095,24 @@ dependencies = [ "libm", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags", +] + +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -1096,6 +1163,25 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_shared", + "serde", +] + +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -1141,6 +1227,35 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6" +[[package]] +name = "postgres-protocol" +version = "0.6.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ee9dd5fe15055d2b6806f4736aa0c9637217074e224bbec46d4041b91bb9491" +dependencies = [ + "base64", + "byteorder", + "bytes", + "fallible-iterator", + "hmac", + "md-5", + "memchr", + "rand 0.9.2", + "sha2", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54b858f82211e84682fecd373f68e1ceae642d8d751a1ebd13f33de6257b3e20" +dependencies = [ + "bytes", + "fallible-iterator", + "postgres-protocol", +] + [[package]] name = "potential_utf" version = "0.1.4" @@ -1210,6 +1325,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" @@ -1223,8 +1344,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", ] [[package]] @@ -1234,7 +1365,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", ] [[package]] @@ -1246,6 +1387,15 @@ dependencies = [ "getrandom 0.2.17", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1264,6 +1414,23 @@ dependencies = [ "bitflags", ] +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + [[package]] name = "rsa" version = "0.9.10" @@ -1277,7 +1444,7 @@ dependencies = [ "num-traits", "pkcs1", "pkcs8", - "rand_core", + "rand_core 0.6.4", "signature", "spki", "subtle", @@ -1396,6 +1563,15 @@ dependencies = [ "digest", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "1.3.0" @@ -1419,9 +1595,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" dependencies = [ "digest", - "rand_core", + "rand_core 0.6.4", ] +[[package]] +name = "siphasher" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" + [[package]] name = "slab" version = "0.4.12" @@ -1580,7 +1762,7 @@ dependencies = [ "memchr", "once_cell", "percent-encoding", - "rand", + "rand 0.8.5", "rsa", "serde", "sha1", @@ -1590,7 +1772,7 @@ dependencies = [ "stringprep", "thiserror", "tracing", - "whoami", + "whoami 1.6.1", ] [[package]] @@ -1618,7 +1800,7 @@ dependencies = [ "md-5", "memchr", "once_cell", - "rand", + "rand 0.8.5", "serde", "serde_json", "sha2", @@ -1627,7 +1809,7 @@ dependencies = [ "stringprep", "thiserror", "tracing", - "whoami", + "whoami 1.6.1", ] [[package]] @@ -1725,6 +1907,15 @@ dependencies = [ "syn", ] +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + [[package]] name = "tinystr" version = "0.8.2" @@ -1778,6 +1969,32 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-postgres" +version = "0.7.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcea47c8f71744367793f16c2db1f11cb859d28f436bdb4ca9193eb1f787ee42" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand 0.9.2", + "socket2", + "tokio", + "tokio-util", + "whoami 2.1.1", +] + [[package]] name = "tokio-stream" version = "0.1.18" @@ -1789,6 +2006,19 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + [[package]] name = "tower" version = "0.5.3" @@ -1847,6 +2077,36 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f30143827ddab0d256fd843b7a66d164e9f271cfa0dde49142c5ca0ca291f1e" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", ] [[package]] @@ -1918,6 +2178,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "vcpkg" version = "0.2.15" @@ -1936,6 +2202,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.2+wasi-0.2.9" @@ -1960,6 +2235,15 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b" +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.114" @@ -2039,6 +2323,16 @@ dependencies = [ "semver", ] +[[package]] +name = "web-sys" +version = "0.3.91" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "854ba17bb104abfb26ba36da9729addc7ce7f06f5c0f90f3c391f8461cca21f9" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "whoami" version = "1.6.1" @@ -2046,7 +2340,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d4a4db5077702ca3015d3d02d74974948aba2ad9e12ab7df718ee64ccd7e97d" dependencies = [ "libredox", - "wasite", + "wasite 0.1.0", +] + +[[package]] +name = "whoami" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6a5b12f9df4f978d2cfdb1bd3bac52433f44393342d7ee9c25f5a1c14c0f45d" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite 1.0.2", + "web-sys", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index b511823..0e827d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,14 +3,20 @@ name = "livevue-rs" version = "0.1.0" edition = "2021" -# Run the todo example with: cargo run --example todo +# Run the todo example with: cargo run --example todo +# Run the pg_replication example with: cargo run --example pg_replication --features pg_replication + +[features] +# Opt-in feature: PostgreSQL logical replication for reactive cache invalidation. +# Adds tokio-postgres and bytes as dependencies. +pg_replication = ["dep:tokio-postgres", "dep:postgres-protocol", "dep:bytes"] [dependencies] axum = { version = "0.7", features = ["macros"] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } serde_json = "1" -sqlx = { version = "0.8", features = ["runtime-tokio", "sqlite"] } +sqlx = { version = "0.8", features = ["runtime-tokio", "sqlite", "postgres"] } maud = { version = "0.26", features = ["axum"] } uuid = { version = "1", features = ["v4", "serde"] } dashmap = "6" @@ -19,6 +25,16 @@ tokio-stream = "0.1" tracing = "0.1.44" futures = "0.3.32" +# Optional: PostgreSQL logical replication +tokio-postgres = { version = "0.7", optional = true } +postgres-protocol = { version = "0.6", optional = true } +bytes = { version = "1", optional = true } [dev-dependencies] -chrono = { version = "0.4", features = ["clock"] } \ No newline at end of file +chrono = { version = "0.4", features = ["clock"] } +tracing-subscriber = { version = "0.3", features = ["env-filter"] } + +[[example]] +name = "pg_replication" +path = "examples/pg_replication/main.rs" +required-features = ["pg_replication"] \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..2efddd9 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,30 @@ +services: + postgres: + image: postgres:16-alpine + container_name: livevue_postgres + restart: unless-stopped + environment: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: livevue + ports: + - "5432:5432" + volumes: + - postgres_data:/var/lib/postgresql/data + # Enable logical replication — required for the pg_replication example. + command: + - postgres + - -c + - wal_level=logical + - -c + - max_replication_slots=10 + - -c + - max_wal_senders=10 + healthcheck: + test: ["CMD-SHELL", "pg_isready -U postgres -d livevue"] + interval: 5s + timeout: 5s + retries: 5 + +volumes: + postgres_data: diff --git a/examples/pg_replication/main.rs b/examples/pg_replication/main.rs new file mode 100644 index 0000000..2f06ce2 --- /dev/null +++ b/examples/pg_replication/main.rs @@ -0,0 +1,387 @@ +//! # pg_replication example +//! +//! A live message board backed by PostgreSQL. +//! +//! Run: +//! ```sh +//! docker compose up -d # start the PG container (see docker-compose.yml) +//! cargo run --example pg_replication --features pg_replication +//! ``` +//! Then open http://localhost:3001 in two browser tabs and watch them stay in sync. +//! +//! ## How it works +//! +//! 1. The render function queries the `messages` table directly via `sqlx::PgPool` +//! and calls `cx.subscribe(pg_table_key("public", "messages"))` to tell the +//! framework "re-render this connection whenever `pg:public.messages` is +//! invalidated". +//! +//! 2. `PgReplicationListener` connects to PostgreSQL over the logical replication +//! protocol, buffers changes per transaction, and on each COMMIT publishes +//! `CacheKey::Channel { key: "pg:public.messages" }` to the store's broadcast +//! channel. +//! +//! 3. The fanout task picks up the event, finds all connections subscribed to that +//! key, and triggers a re-render for each — sending fresh HTML over SSE. +//! +//! No polling. No websocket boilerplate. Just PostgreSQL WAL + SSE. + +use std::sync::Arc; +use std::time::Duration; + +use axum::{ + extract::{Json, State}, + http::StatusCode, + response::{Html, IntoResponse}, + routing::{get, post}, + Router, +}; +use maud::{html, Markup, DOCTYPE}; +use serde::Deserialize; +use sqlx::FromRow; +use uuid::Uuid; + +use livevue_rs::{ + pg_replication::{pg_table_key, PgReplicationConfig, PgReplicationListener}, + server::{ingest_signals, ActionParams, AppState, SharedState}, + spawn_fanout, RenderContext, Store, +}; + +// --------------------------------------------------------------------------- +// Domain types +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, FromRow)] +struct Message { + id: i64, + body: String, +} + +// --------------------------------------------------------------------------- +// Database helpers +// --------------------------------------------------------------------------- + +async fn setup_db(pool: &sqlx::PgPool) -> anyhow::Result<()> { + // Enable logical replication for this table (requires superuser or + // rds_superuser on managed PG). + sqlx::query( + "CREATE TABLE IF NOT EXISTS messages ( + id BIGSERIAL PRIMARY KEY, + body TEXT NOT NULL + )", + ) + .execute(pool) + .await?; + + // Create the publication if it doesn't exist. + // `IF NOT EXISTS` was added in PG 14 — fall back gracefully on older versions. + let exists: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_publication WHERE pubname = 'livevue_pub')", + ) + .fetch_one(pool) + .await?; + + if !exists { + sqlx::query("CREATE PUBLICATION livevue_pub FOR TABLE messages") + .execute(pool) + .await?; + println!("Created publication livevue_pub"); + } + + Ok(()) +} + +// --------------------------------------------------------------------------- +// Rendering +// --------------------------------------------------------------------------- + +/// Shared render state — anything the render function needs beyond the store. +#[derive(Clone)] +struct RenderState { + pg: sqlx::PgPool, +} + +async fn render_page(cx: &mut RenderContext, rs: &RenderState) -> anyhow::Result { + // Subscribe to PG-sourced invalidation events for the messages table. + // The fanout task will re-render this connection whenever the WAL emitter + // fires a `CacheKey::Channel { key: "pg:public.messages" }` event. + cx.subscribe(pg_table_key("public", "messages")); + + let messages: Vec = sqlx::query_as("SELECT id, body FROM messages ORDER BY id DESC") + .fetch_all(&rs.pg) + .await?; + + let new_msg = cx.signal("newMsg", livevue_rs::global()); + + Ok(html! { + div #app { + h1 { "Live Message Board" } + p.subtitle { + "Messages are stored in PostgreSQL. " + "Add one and watch every open tab update in real time via WAL replication." + } + + // ── New message form ───────────────────────────────────────── + div.compose { + input + type="text" + placeholder="Type a message…" + data-bind=(new_msg.name) + data-on:keydown="{if (event.key === 'Enter') $post('/action/post_message')}" + ; + button + data-on:click="$post('/action/post_message')" + data-attr=(format!("{{disabled: {}.trim() === ''}}", new_msg.val)) + { "Send" } + } + + // ── Message list ───────────────────────────────────────────── + @if messages.is_empty() { + p.empty { "No messages yet. Be the first!" } + } @else { + ul.messages { + @for msg in &messages { + li { + span.id { "#" (msg.id) } + span.body { (msg.body) } + } + } + } + } + } + }) +} + +/// Full HTML shell for the initial page load. +fn page_shell(conn_id: Uuid, signals_json: &str, inner: &Markup) -> String { + html! { + (DOCTYPE) + html lang="en" { + head { + meta charset="utf-8"; + meta name="viewport" content="width=device-width, initial-scale=1"; + title { "pg_replication example — livevue-rs" } + style { + (r#" + body { font-family: system-ui, sans-serif; max-width: 640px; margin: 2rem auto; padding: 0 1rem; } + h1 { margin-bottom: 0.25rem; } + .subtitle { color: #666; font-size: 0.9rem; margin-top: 0; } + .compose { display: flex; gap: 0.5rem; margin: 1.5rem 0; } + .compose input { flex: 1; padding: 0.5rem; border: 1px solid #ccc; border-radius: 4px; font-size: 1rem; } + .compose button { padding: 0.5rem 1rem; background: #0070f3; color: white; border: none; border-radius: 4px; cursor: pointer; font-size: 1rem; } + .compose button:disabled { background: #aaa; cursor: default; } + ul.messages { list-style: none; padding: 0; margin: 0; } + ul.messages li { padding: 0.75rem; border-bottom: 1px solid #eee; display: flex; gap: 0.75rem; } + .id { color: #999; font-size: 0.85rem; min-width: 2.5rem; } + .empty { color: #999; } + "#) + } + // DataStar SDK (loaded from CDN for the example) + script + type="module" + src="https://cdn.jsdelivr.net/gh/starfederation/datastar@v1.0.0-beta.11/bundles/datastar.js" + {} + } + body { + div + data-signals=(signals_json) + data-on:load=(format!("@get('/sse?conn={conn_id}')")) + { + (inner) + } + } + } + } + .into_string() +} + +// --------------------------------------------------------------------------- +// HTTP handlers +// --------------------------------------------------------------------------- + +async fn index_handler(State(state): State<(SharedState, RenderState)>) -> impl IntoResponse { + let (app_state, rs) = &state; + let conn_id = Uuid::new_v4(); + let signals_json = r#"{"newMsg":""}"#; + app_state.connections.insert(conn_id, { + // We need a sender for the connection, but the SSE handler will replace + // it. Use a dummy channel that immediately drops. + let (tx, _rx) = tokio::sync::mpsc::channel(1); + tx + }); + app_state + .connections + .update_signals(&conn_id, signals_json.to_string()); + + let mut cx = + RenderContext::new(conn_id, livevue_rs::SignalStore::from_json(serde_json::json!({"newMsg": ""})), Default::default(), app_state.store.clone()); + + let inner = render_page(&mut cx, rs).await.unwrap_or_else(|e| { + html! { p { "Render error: " (e) } } + }); + + Html(page_shell(conn_id, signals_json, &inner)) +} + +// SSE handler is provided by the framework; it needs SharedState only. +// We wrap it to adapt the two-tuple state. +async fn sse_adapter( + State(state): State<(SharedState, RenderState)>, + query: axum::extract::Query, +) -> impl IntoResponse { + livevue_rs::sse_handler(State(state.0), query).await +} + +#[derive(Deserialize)] +struct PostMessageBody { + #[serde(rename = "newMsg")] + new_msg: Option, +} + +async fn post_message_handler( + State(state): State<(SharedState, RenderState)>, + axum::extract::Query(params): axum::extract::Query, + Json(body): Json, +) -> impl IntoResponse { + let signals = ingest_signals(&state.0, params.conn, body); + let msg: String = signals + .get("newMsg") + .and_then(|v| v.as_str().map(|s| s.trim().to_string())) + .unwrap_or_default(); + + if msg.is_empty() { + return StatusCode::BAD_REQUEST; + } + + if let Err(e) = sqlx::query("INSERT INTO messages (body) VALUES ($1)") + .bind(&msg) + .execute(&state.1.pg) + .await + { + tracing::error!("failed to insert message: {e}"); + return StatusCode::INTERNAL_SERVER_ERROR; + } + + // Note: we do NOT call store.publish() here — the PgReplicationListener + // will detect the INSERT through the WAL stream and fire the invalidation + // automatically. This demonstrates the key value of the module: mutations + // from *any* source (other services, psql, migrations) are reflected + // without any extra wiring. + + StatusCode::OK +} + +fn router(app_state: SharedState, render_state: RenderState) -> Router { + let combined = (app_state, render_state); + Router::new() + .route("/", get(index_handler)) + .route("/sse", get(sse_adapter)) + .route("/action/post_message", post(post_message_handler)) + .with_state(combined) +} + +// --------------------------------------------------------------------------- +// main +// --------------------------------------------------------------------------- + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::from_default_env() + .add_directive("livevue_rs=debug".parse().unwrap()) + .add_directive("pg_replication=debug".parse().unwrap()), + ) + .init(); + + // ----------------------------------------------------------------------- + // Database setup + // ----------------------------------------------------------------------- + + let database_url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:postgres@localhost:5432/livevue".into()); + + let pg = sqlx::PgPool::connect(&database_url).await?; + setup_db(&pg).await?; + + // ----------------------------------------------------------------------- + // Framework store (SQLite in-memory, unused for queries here) + // ----------------------------------------------------------------------- + + // The Store provides the broadcast channel and KV store that the framework + // needs. The embedded SQLite pool is not used in this example — queries go + // directly to PostgreSQL via the captured `pg` pool in the render closure. + let sqlite = sqlx::SqlitePool::connect("sqlite::memory:").await?; + let store = Store::new(sqlite); + + // ----------------------------------------------------------------------- + // Render function + // ----------------------------------------------------------------------- + + let render_state = RenderState { pg: pg.clone() }; + let rs_clone = render_state.clone(); + + let render_fn: livevue_rs::RenderFn = Arc::new(move |mut cx| { + let rs = rs_clone.clone(); + Box::pin(async move { + let signals_json = cx.signals.to_json_string(); + let conn_id = cx.connection_id; + + let inner = render_page(&mut cx, &rs).await?; + let keys = cx.take_subscription_keys(); + + // Wrap in the full page shell so DataStar can morph the entire document. + let html = page_shell(conn_id, &signals_json, &inner); + Ok((html, keys)) + }) + }); + + // ----------------------------------------------------------------------- + // Application state + fanout + // ----------------------------------------------------------------------- + + let state = AppState::new(store.clone(), render_fn); + spawn_fanout(state.clone()); + + // ----------------------------------------------------------------------- + // PG replication listener + // ----------------------------------------------------------------------- + + // The replication protocol requires a dedicated connection. + + let replication_store = store.clone(); + tokio::spawn(async move { + loop { + // Rebuild the config each iteration — PgReplicationConfig is not + // Clone, and start() takes ownership, so we reconstruct from the + // env-var values cached in the outer scope. + let c = PgReplicationConfig { + host: std::env::var("PG_HOST").unwrap_or_else(|_| "localhost".into()), + port: std::env::var("PG_PORT").ok() + .and_then(|p| p.parse().ok()).unwrap_or(5432), + user: std::env::var("PG_USER").unwrap_or_else(|_| "postgres".into()), + password: std::env::var("PG_PASSWORD").unwrap_or_else(|_| "postgres".into()), + database: std::env::var("PG_DATABASE").unwrap_or_else(|_| "livevue".into()), + slot_name: "livevue_slot".into(), + publication_name: "livevue_pub".into(), + start_lsn: None, + }; + if let Err(e) = PgReplicationListener::new(c, replication_store.clone()).start().await { + tracing::error!("pg replication listener error: {e:#}"); + } + tracing::info!("pg replication listener reconnecting in 5 s…"); + tokio::time::sleep(Duration::from_secs(5)).await; + } + }); + + // ----------------------------------------------------------------------- + // HTTP server + // ----------------------------------------------------------------------- + + let app = router(state, render_state); + let listener_tcp = tokio::net::TcpListener::bind("0.0.0.0:3001").await?; + println!("pg_replication example listening on http://localhost:3001"); + axum::serve(listener_tcp, app).await?; + + Ok(()) +} diff --git a/src/context.rs b/src/context.rs index c9c220d..00cc17e 100644 --- a/src/context.rs +++ b/src/context.rs @@ -197,6 +197,30 @@ impl RenderContext { self.store.kv.get(&key) } + // ----------------------------------------------------------------------- + // Manual subscriptions + // ----------------------------------------------------------------------- + + /// Manually register a subscription key without executing a query. + /// + /// Use this to subscribe to externally-sourced invalidation events that + /// don't go through `cx.run()` — for example, events published by the + /// PostgreSQL replication module: + /// + /// ```ignore + /// use livevue_rs::pg_replication::pg_table_key; + /// + /// cx.subscribe(pg_table_key("public", "posts")); + /// let posts = sqlx::query_as::<_, Post>("SELECT id, body FROM posts") + /// .fetch_all(&pg_pool) + /// .await?; + /// ``` + pub fn subscribe(&mut self, key: CacheKey) { + if key != CacheKey::None { + self.subscription_keys.push(key); + } + } + // ----------------------------------------------------------------------- // Post-render access // ----------------------------------------------------------------------- diff --git a/src/lib.rs b/src/lib.rs index a1bb017..7a6d34e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,15 @@ pub mod connection; pub mod sse; pub mod server; +/// PostgreSQL logical replication module. +/// +/// Enabled via the `pg_replication` Cargo feature: +/// ```toml +/// livevue-rs = { features = ["pg_replication"] } +/// ``` +#[cfg(feature = "pg_replication")] +pub mod pg_replication; + pub use signal::{Signal, SignalStore, SignalOpts, mangle, client_only, global, data_signals}; pub use query::{CacheKey, Query}; pub use store::{Store, KVStore}; diff --git a/src/pg_replication/emitter.rs b/src/pg_replication/emitter.rs new file mode 100644 index 0000000..4309b16 --- /dev/null +++ b/src/pg_replication/emitter.rs @@ -0,0 +1,311 @@ +use super::wal_parser::{ColumnValue, LogicalMessage, MessageParser, RelationInfo, TupleData}; +use crate::query::CacheKey; +use bytes::BytesMut; +use std::collections::HashSet; +use tokio::sync::broadcast; + +// --------------------------------------------------------------------------- +// Per-transaction buffer +// --------------------------------------------------------------------------- + +/// Accumulates the set of invalidation keys for one PostgreSQL transaction. +/// +/// We use a `HashSet` for table-level keys so that a transaction inserting +/// 10 000 rows into one table still produces exactly **one** +/// `CacheKey::Channel { key: "pg:public.todos" }` event. Row-level keys are +/// collected separately for callers that need fine-grained precision. +struct TxBuffer { + xid: u32, + /// Deduplicated `(schema, table)` pairs touched in this transaction. + table_keys: HashSet<(String, String)>, + /// `(schema, table, id)` tuples for changes with a detectable integer PK. + row_keys: Vec<(String, String, i64)>, +} + +impl TxBuffer { + fn new(xid: u32) -> Self { + Self { + xid, + table_keys: HashSet::new(), + row_keys: Vec::new(), + } + } + + fn record_table(&mut self, schema: &str, table: &str) { + self.table_keys.insert((schema.to_owned(), table.to_owned())); + } + + fn record_row(&mut self, schema: &str, table: &str, id: i64) { + self.row_keys.push((schema.to_owned(), table.to_owned(), id)); + } +} + +// --------------------------------------------------------------------------- +// WalEmitter +// --------------------------------------------------------------------------- + +/// Processes pgoutput logical replication messages and publishes +/// [`CacheKey`] invalidation events to the framework's broadcast channel. +/// +/// ## Transaction atomicity +/// +/// Changes are buffered per XID. On `COMMIT`, all invalidation events for +/// that transaction are published in a tight loop so the fanout task receives +/// them as a burst. The transaction's WAL data is **never** partially visible: +/// the emitter waits for the `Commit` message before publishing anything. +/// +/// ## Event granularity +/// +/// For every transaction two kinds of events are emitted: +/// +/// * **Table-level** (`"pg:{schema}.{table}"`) — deduplicated; at most one per +/// distinct table touched. Components that track whole-table queries should +/// subscribe with [`pg_table_key`](super::pg_table_key). +/// +/// * **Row-level** (`"pg:{schema}.{table}:{id}"`) — emitted when the change +/// has a single-column integer replica-identity key. Components that track +/// individual rows can subscribe with [`pg_row_key`](super::pg_row_key). +/// +/// ## Usage +/// +/// ```ignore +/// let emitter = WalEmitter::new(store.events.clone()); +/// // Inside the replication loop: +/// emitter.process_raw(&wal_bytes).await?; +/// ``` +pub struct WalEmitter { + parser: MessageParser, + events: broadcast::Sender, + /// XID of the currently open transaction (pgoutput v1 is sequential). + current_xid: Option, + /// Buffer for the current transaction; `None` when outside a transaction. + buffer: Option, +} + +impl WalEmitter { + /// Create a new emitter that publishes to `events`. + /// + /// Pass `store.events.clone()` so the emitter integrates directly with the + /// framework's fanout task. + pub fn new(events: broadcast::Sender) -> Self { + Self { + parser: MessageParser::new(), + events, + current_xid: None, + buffer: None, + } + } + + // ----------------------------------------------------------------------- + // Public entry point + // ----------------------------------------------------------------------- + + /// Parse and process all pgoutput messages in a raw XLogData payload. + /// + /// Call this for every `XLogData` message received from the replication + /// stream. The method is `&mut self` because the parser maintains relation + /// metadata across calls. + pub async fn process_raw(&mut self, data: &[u8]) -> Result<(), String> { + let mut buf = BytesMut::from(data); + while !buf.is_empty() { + match self.parser.parse_message(&mut buf)? { + Some(msg) => self.process_message(msg)?, + None => break, + } + } + Ok(()) + } + + // ----------------------------------------------------------------------- + // Message dispatch + // ----------------------------------------------------------------------- + + fn process_message(&mut self, msg: LogicalMessage) -> Result<(), String> { + match msg { + LogicalMessage::Begin { xid, .. } => { + self.current_xid = Some(xid); + self.buffer = Some(TxBuffer::new(xid)); + } + + LogicalMessage::Commit { .. } => { + if let Some(buf) = self.buffer.take() { + self.emit(buf); + } + self.current_xid = None; + } + + LogicalMessage::Insert { rel_oid, new_tuple } => { + self.record_change(rel_oid, &new_tuple); + } + + LogicalMessage::Update { rel_oid, new_tuple, key_tuple, .. } => { + // Prefer the key_tuple for row-level id extraction if available + // (it holds the *old* key when the PK is being changed). + let id_source = key_tuple.as_ref().unwrap_or(&new_tuple); + let id = self.extract_pk_id(rel_oid, id_source); + self.record_change_with_id(rel_oid, &new_tuple, id); + } + + LogicalMessage::Delete { rel_oid, key_tuple, old_tuple } => { + let tuple = key_tuple.as_ref().or(old_tuple.as_ref()); + let id = tuple.and_then(|t| self.extract_pk_id(rel_oid, t)); + // We still need the relation name; use an empty TupleData as placeholder. + if let Some(info) = self.relation_info(rel_oid) { + let schema = info.0; + let table = info.1; + if let Some(buf) = &mut self.buffer { + buf.record_table(&schema, &table); + if let Some(id) = id { + buf.record_row(&schema, &table, id); + } + } + } + } + + LogicalMessage::Truncate { rel_oids, .. } => { + let pairs: Vec<(String, String)> = rel_oids + .iter() + .filter_map(|&oid| self.relation_info(oid)) + .collect(); + if let Some(buf) = &mut self.buffer { + for (schema, table) in pairs { + buf.record_table(&schema, &table); + } + } + } + + // Relation messages are consumed by the parser (internal state). + // Type and Message are not relevant for cache invalidation. + LogicalMessage::Relation(_) + | LogicalMessage::Type { .. } + | LogicalMessage::Message { .. } => {} + } + + Ok(()) + } + + // ----------------------------------------------------------------------- + // Helpers + // ----------------------------------------------------------------------- + + /// Record a table + optional row key for a DML change. + fn record_change(&mut self, rel_oid: u32, tuple: &TupleData) { + let id = self.extract_pk_id(rel_oid, tuple); + self.record_change_with_id(rel_oid, tuple, id); + } + + fn record_change_with_id(&mut self, rel_oid: u32, _tuple: &TupleData, id: Option) { + if let Some(info) = self.relation_info(rel_oid) { + let (schema, table) = info; + if let Some(buf) = &mut self.buffer { + buf.record_table(&schema, &table); + if let Some(id) = id { + buf.record_row(&schema, &table, id); + } + } + } + } + + /// Resolve a relation OID to `(schema, table)`. Clones so the parser + /// borrow is released before we mutably borrow `self.buffer`. + fn relation_info(&self, rel_oid: u32) -> Option<(String, String)> { + self.parser + .get_relation(rel_oid) + .map(|r| (r.namespace.clone(), r.name.clone())) + } + + /// Try to extract a single-column integer primary key from `tuple`. + /// + /// Returns `None` for composite keys, non-integer keys, or when the + /// relation has no replica-identity columns flagged (`flags & 1`). + fn extract_pk_id(&self, rel_oid: u32, tuple: &TupleData) -> Option { + let rel = self.parser.get_relation(rel_oid)?; + extract_single_int_pk(rel, tuple) + } + + // ----------------------------------------------------------------------- + // Emission + // ----------------------------------------------------------------------- + + /// Publish all invalidation events for a committed transaction. + fn emit(&self, buf: TxBuffer) { + tracing::debug!( + xid = buf.xid, + tables = buf.table_keys.len(), + rows = buf.row_keys.len(), + "pg_replication: emitting cache invalidation" + ); + + // Table-level events (deduplicated by HashSet). + for (schema, table) in &buf.table_keys { + self.events + .send(CacheKey::Channel { + key: format!("pg:{schema}.{table}"), + }) + .ok(); // ok() = ignore when no receivers are subscribed + } + + // Row-level events. + for (schema, table, id) in &buf.row_keys { + self.events + .send(CacheKey::Channel { + key: format!("pg:{schema}.{table}:{id}"), + }) + .ok(); + } + } +} + +// --------------------------------------------------------------------------- +// Row-level PK extraction +// --------------------------------------------------------------------------- + +/// Extract an integer PK only when the relation has exactly one column +/// flagged as part of the replica identity key and its value is parseable +/// as `i64`. +fn extract_single_int_pk(rel: &RelationInfo, tuple: &TupleData) -> Option { + let key_cols: Vec = rel + .columns + .iter() + .enumerate() + .filter(|(_, col)| col.flags & 1 != 0) + .map(|(i, _)| i) + .collect(); + + if key_cols.len() != 1 { + return None; + } + + let idx = key_cols[0]; + match tuple.values.get(idx)? { + ColumnValue::Text(s) => s.parse::().ok(), + _ => None, + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn no_receivers_does_not_panic() { + let (tx, _rx) = broadcast::channel(16); + drop(_rx); // no receivers + let mut emitter = WalEmitter::new(tx); + // Publishing to a channel with no receivers should silently succeed. + emitter + .emit(TxBuffer { xid: 1, table_keys: Default::default(), row_keys: vec![] }); + } + + #[test] + fn outside_transaction_dml_is_ignored() { + let (tx, _rx) = broadcast::channel(16); + let mut emitter = WalEmitter::new(tx); + // No Begin sent → buffer is None → Insert is silently dropped. + assert!(emitter.buffer.is_none()); + } +} diff --git a/src/pg_replication/mod.rs b/src/pg_replication/mod.rs new file mode 100644 index 0000000..a39eb87 --- /dev/null +++ b/src/pg_replication/mod.rs @@ -0,0 +1,265 @@ +//! PostgreSQL logical replication module for livevue-rs. +//! +//! Connects to a PostgreSQL `pgoutput` logical replication stream, buffers +//! changes per transaction, and publishes [`CacheKey`] invalidation events +//! to the framework's broadcast channel on every committed transaction. +//! +//! # Quick start +//! +//! ```ignore +//! use livevue_rs::pg_replication::{PgReplicationConfig, PgReplicationListener, pg_table_key}; +//! +//! let config = PgReplicationConfig { +//! host: "localhost".into(), +//! port: 5432, +//! user: "postgres".into(), +//! password: "postgres".into(), +//! database: "mydb".into(), +//! slot_name: "livevue_slot".into(), +//! publication_name: "livevue_pub".into(), +//! start_lsn: None, +//! }; +//! +//! let listener = PgReplicationListener::new(config, store.clone()); +//! tokio::spawn(async move { +//! if let Err(e) = listener.start().await { +//! tracing::error!("replication listener exited: {e}"); +//! } +//! }); +//! ``` +//! +//! # CacheKey convention +//! +//! | Granularity | Channel key format | Example | +//! |-------------|------------------------------|-----------------------------| +//! | Table-level | `"pg:{schema}.{table}"` | `"pg:public.posts"` | +//! | Row-level | `"pg:{schema}.{table}:{id}"` | `"pg:public.posts:42"` | +//! +//! Use [`pg_table_key`] / [`pg_row_key`] in `cx.subscribe()` in render functions. +//! +//! # Feature flag +//! +//! Enable with `features = ["pg_replication"]` in your `Cargo.toml`. + +mod wal_parser; +mod emitter; +mod proto; + +pub use wal_parser::{ + ColumnInfo, ColumnValue, LogicalMessage, MessageParser, RelationInfo, TupleData, +}; +pub use emitter::WalEmitter; + +use crate::query::CacheKey; +use crate::store::Store; +use bytes::{BufMut, Bytes, BytesMut}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +// --------------------------------------------------------------------------- +// Key helpers +// --------------------------------------------------------------------------- + +/// [`CacheKey`] matching table-level events from the PG replication module. +/// +/// Call `cx.subscribe(pg_table_key("public", "messages"))` in a render +/// function to re-render whenever any row in `messages` changes. +pub fn pg_table_key(schema: &str, table: &str) -> CacheKey { + CacheKey::Channel { key: format!("pg:{schema}.{table}") } +} + +/// [`CacheKey`] matching row-level events from the PG replication module. +/// +/// Row-level events are only emitted when the row has a single-column +/// integer replica-identity key. +pub fn pg_row_key(schema: &str, table: &str, id: i64) -> CacheKey { + CacheKey::Channel { key: format!("pg:{schema}.{table}:{id}") } +} + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/// Configuration for a [`PgReplicationListener`]. +pub struct PgReplicationConfig { + pub host: String, + pub port: u16, + pub user: String, + pub password: String, + pub database: String, + + /// Logical replication slot name. Auto-created with the `pgoutput` + /// plugin if it does not already exist. + pub slot_name: String, + + /// Publication name. Must be created in PostgreSQL first: + /// ```sql + /// CREATE PUBLICATION livevue_pub FOR ALL TABLES; + /// ``` + pub publication_name: String, + + /// LSN to start from. `None` = start from the slot's confirmed LSN + /// (only future changes). + pub start_lsn: Option, +} + +// --------------------------------------------------------------------------- +// Listener +// --------------------------------------------------------------------------- + +/// Streams PostgreSQL WAL changes and publishes [`CacheKey`] events to the +/// framework's broadcast channel. +/// +/// Spawn as a long-running background task after calling `AppState::new` and +/// `spawn_fanout`. Wrap in a retry loop for automatic reconnection. +pub struct PgReplicationListener { + config: PgReplicationConfig, + store: Arc, +} + +impl PgReplicationListener { + pub fn new(config: PgReplicationConfig, store: Arc) -> Self { + Self { config, store } + } + + /// Connect and stream WAL until an error or stream close. + pub async fn start(self) -> anyhow::Result<()> { + use anyhow::Context as _; + + // --- connect in replication mode --- + let mut conn = proto::RawPgConn::connect( + &self.config.host, + self.config.port, + &self.config.user, + &self.config.password, + &self.config.database, + ) + .await + .context("connect to PostgreSQL (replication mode)")?; + + // --- ensure replication slot exists --- + let slot = &self.config.slot_name; + let create_sql = format!( + "CREATE_REPLICATION_SLOT {slot} LOGICAL pgoutput NOEXPORT_SNAPSHOT" + ); + match conn.simple_query(&create_sql).await { + Ok(_) => tracing::info!(slot, "pg_replication: created replication slot"), + Err(e) => { + let msg = e.to_string(); + if msg.contains("42710") || msg.contains("already exists") { + tracing::debug!(slot, "pg_replication: slot already exists"); + } else { + return Err(e).context("create replication slot"); + } + } + } + + // --- start logical replication --- + let lsn = self.config.start_lsn.unwrap_or(0); + let lsn_str = format!("{:X}/{:X}", (lsn >> 32) as u32, lsn as u32); + let pub_name = &self.config.publication_name; + let start_sql = format!( + "START_REPLICATION SLOT {slot} LOGICAL {lsn_str} \ + (proto_version '1', publication_names '{pub_name}')" + ); + + conn.start_replication(&start_sql) + .await + .context("START_REPLICATION")?; + + tracing::info!( + slot, publication = pub_name, lsn = lsn_str, + "pg_replication: replication stream started" + ); + + // --- stream processing loop --- + let mut emitter = WalEmitter::new(self.store.events.clone()); + let mut last_lsn: u64 = lsn; + let mut last_status_at = Instant::now(); + const STATUS_INTERVAL: Duration = Duration::from_secs(10); + + loop { + let frame = tokio::time::timeout( + Duration::from_secs(1), + conn.read_copy_data(), + ) + .await; + + match frame { + Err(_timeout) => {} + + Ok(Err(e)) => return Err(e).context("replication stream read"), + + Ok(Ok(None)) => { + tracing::info!("pg_replication: replication stream ended"); + break; + } + + Ok(Ok(Some(data))) => { + if data.is_empty() { continue; } + match data[0] { + // XLogData: 'w' + u64 wal_start(8) + u64 wal_end(8) + i64 time(8) + WAL + b'w' if data.len() >= 25 => { + last_lsn = u64::from_be_bytes( + data[9..17].try_into().expect("8 byte slice"), + ); + if let Err(e) = emitter.process_raw(&data[25..]).await { + tracing::warn!("pg_replication: WAL parse warning: {e}"); + } + } + // PrimaryKeepalive: 'k' + u64 wal_end(8) + i64 time(8) + u8 reply_req + b'k' if data.len() >= 18 => { + if data[17] != 0 { + conn.send_standby_status(last_lsn) + .await + .context("send standby status (keepalive reply)")?; + last_status_at = Instant::now(); + } + } + other => { + tracing::debug!( + "pg_replication: unknown replication msg 0x{:02x}", other + ); + } + } + } + } + + if last_status_at.elapsed() >= STATUS_INTERVAL { + conn.send_standby_status(last_lsn) + .await + .context("send periodic standby status")?; + last_status_at = Instant::now(); + } + } + + Ok(()) + } +} + +// --------------------------------------------------------------------------- +// Internal helpers +// --------------------------------------------------------------------------- + +/// StandbyStatusUpdate payload for acknowledging processed WAL. +pub(crate) fn standby_status_payload(lsn: u64) -> Bytes { + let ack = lsn.saturating_add(1); + let mut buf = BytesMut::with_capacity(34); + buf.put_u8(b'r'); + buf.put_u64(ack); // write LSN + buf.put_u64(ack); // flush LSN + buf.put_u64(ack); // apply LSN + buf.put_i64(pg_epoch_micros()); + buf.put_u8(0); // no reply requested + buf.freeze() +} + +/// Current time as microseconds since the PostgreSQL epoch (2000-01-01 UTC). +pub(crate) fn pg_epoch_micros() -> i64 { + use std::time::{SystemTime, UNIX_EPOCH}; + let micros = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_micros() as i64; + micros - 946_684_800_000_000 +} diff --git a/src/pg_replication/proto.rs b/src/pg_replication/proto.rs new file mode 100644 index 0000000..46c867d --- /dev/null +++ b/src/pg_replication/proto.rs @@ -0,0 +1,341 @@ +//! Raw PostgreSQL wire-protocol connection for logical replication. +//! +//! `tokio-postgres` 0.7 does not expose a `copy_both_simple` API suitable for +//! the replication stream. This module implements the minimum subset of the +//! PostgreSQL wire protocol needed to: +//! +//! 1. Connect with `replication=database` startup parameter. +//! 2. Perform MD5 or cleartext password authentication. +//! 3. Send simple queries (for `CREATE_REPLICATION_SLOT` and `START_REPLICATION`). +//! 4. Enter `COPY BOTH` mode and stream `CopyData` frames. +//! 5. Send `StandbyStatusUpdate` CopyData frames back to the server. +//! +//! `postgres-protocol` is used for building most frontend messages and for +//! MD5 hash computation. The backend parser is bypassed for `CopyBothResponse` +//! (tag `'W'` / 0x57) which is not present in that crate's public API. + +use super::standby_status_payload; +use anyhow::{bail, Context as _}; +use bytes::{Buf, Bytes, BytesMut}; +use postgres_protocol::authentication::md5_hash; +use postgres_protocol::message::frontend; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; + +// --------------------------------------------------------------------------- +// RawPgConn +// --------------------------------------------------------------------------- + +/// A raw PostgreSQL connection suitable for replication streaming. +pub struct RawPgConn { + stream: TcpStream, + /// Accumulates bytes read from the server until a complete frame arrives. + read_buf: BytesMut, +} + +impl RawPgConn { + // ----------------------------------------------------------------------- + // Constructor + // ----------------------------------------------------------------------- + + /// Connect to a PostgreSQL server in **logical replication** mode. + /// + /// Sends a startup message with `replication=database` so the backend + /// enters walsender mode. Handles MD5 and cleartext password auth. + pub async fn connect( + host: &str, + port: u16, + user: &str, + password: &str, + database: &str, + ) -> anyhow::Result { + let addr = format!("{host}:{port}"); + let stream = TcpStream::connect(&addr) + .await + .with_context(|| format!("TCP connect to {addr}"))?; + + let mut conn = RawPgConn { + stream, + read_buf: BytesMut::with_capacity(4096), + }; + + // --- Startup message ----------------------------------------------- + // Parameters include `replication=database` which instructs PostgreSQL + // to accept replication commands on this connection. + let mut startup_buf = BytesMut::new(); + frontend::startup_message( + [ + ("user", user), + ("database", database), + ("replication", "database"), + ("application_name", "livevue-rs"), + ] + .iter() + .copied(), + &mut startup_buf, + )?; + conn.stream.write_all(&startup_buf).await?; + conn.stream.flush().await?; + + // --- Authentication loop ------------------------------------------- + loop { + let (tag, body) = conn.read_frame().await?; + match tag { + b'R' => { + // Authentication request + conn.handle_auth_request(&body, user, password).await?; + } + b'K' => { + // BackendKeyData — ignore + } + b'S' => { + // ParameterStatus — ignore + } + b'Z' => { + // ReadyForQuery — handshake complete + break; + } + b'E' => { + let msg = parse_error_message(&body); + bail!("PostgreSQL startup error: {msg}"); + } + b'N' => { + // NoticeResponse — ignore + } + other => { + tracing::debug!("startup: unexpected tag 0x{:02x}", other); + } + } + } + + Ok(conn) + } + + // ----------------------------------------------------------------------- + // Simple query (used for CREATE_REPLICATION_SLOT etc.) + // ----------------------------------------------------------------------- + + /// Send a simple query and drain responses until `ReadyForQuery`. + /// + /// Returns `Err` if the server sends an `ErrorResponse`. + pub async fn simple_query(&mut self, query: &str) -> anyhow::Result<()> { + let mut buf = BytesMut::new(); + frontend::query(query, &mut buf)?; + self.stream.write_all(&buf).await?; + self.stream.flush().await?; + + loop { + let (tag, body) = self.read_frame().await?; + match tag { + b'Z' => break, // ReadyForQuery + b'E' => { + let msg = parse_error_message(&body); + bail!("query error: {msg}"); + } + _ => {} // CommandComplete, RowDescription, DataRow, NoticeResponse … + } + } + Ok(()) + } + + // ----------------------------------------------------------------------- + // Enter COPY BOTH mode for replication streaming + // ----------------------------------------------------------------------- + + /// Send `START_REPLICATION` and absorb the `CopyBothResponse`. + /// + /// After this call the connection is in COPY BOTH mode. Use + /// [`read_copy_data`] / [`send_standby_status`] to exchange messages. + pub async fn start_replication(&mut self, query: &str) -> anyhow::Result<()> { + let mut buf = BytesMut::new(); + frontend::query(query, &mut buf)?; + self.stream.write_all(&buf).await?; + self.stream.flush().await?; + + // Expect CopyBothResponse ('W') followed by the replication stream. + // postgres-protocol doesn't parse 'W', so we handle it manually. + loop { + let (tag, _body) = self.read_frame().await?; + match tag { + b'W' => { + // CopyBothResponse — body contains column-format info we + // don't need for the replication stream. Just break out. + break; + } + b'E' => { + let msg = parse_error_message(&_body); + bail!("START_REPLICATION error: {msg}"); + } + b'N' => {} // NoticeResponse — skip + other => { + tracing::debug!("start_replication: unexpected tag 0x{:02x}", other); + } + } + } + Ok(()) + } + + // ----------------------------------------------------------------------- + // Replication stream I/O + // ----------------------------------------------------------------------- + + /// Read one `CopyData` payload from the replication stream. + /// + /// Returns: + /// - `Ok(Some(bytes))` — a CopyData payload (XLogData or PrimaryKeepalive). + /// - `Ok(None)` — the server sent `CopyDone` or `ErrorResponse` (stream end). + /// - `Err(_)` — I/O or protocol error. + pub async fn read_copy_data(&mut self) -> anyhow::Result> { + let (tag, body) = self.read_frame().await?; + match tag { + b'd' => Ok(Some(body)), + b'c' => Ok(None), // CopyDone + b'E' => { + let msg = parse_error_message(&body); + bail!("replication error: {msg}"); + } + other => { + tracing::debug!("copy stream: unexpected tag 0x{:02x}", other); + Ok(Some(Bytes::new())) + } + } + } + + /// Send a `StandbyStatusUpdate` CopyData frame to acknowledge `lsn`. + pub async fn send_standby_status(&mut self, lsn: u64) -> anyhow::Result<()> { + let payload = standby_status_payload(lsn); + let mut buf = BytesMut::new(); + frontend::CopyData::new(&payload[..])?.write(&mut buf); + self.stream.write_all(&buf).await?; + self.stream.flush().await?; + Ok(()) + } + + // ----------------------------------------------------------------------- + // Low-level frame I/O + // ----------------------------------------------------------------------- + + /// Read one complete PostgreSQL message frame. + /// + /// Frame format: `tag (u8) | length (i32, BE, includes itself) | body`. + /// Returns `(tag, body_bytes)`. + async fn read_frame(&mut self) -> anyhow::Result<(u8, Bytes)> { + // Ensure we have at least 5 bytes (1 tag + 4 length). + while self.read_buf.len() < 5 { + let n = self.stream.read_buf(&mut self.read_buf).await?; + if n == 0 { + bail!("server closed connection unexpectedly"); + } + } + + let tag = self.read_buf[0]; + // Length field is i32 big-endian, includes the 4-length bytes but NOT + // the tag byte. Total frame length = 1 (tag) + length. + let body_len = (i32::from_be_bytes( + self.read_buf[1..5].try_into().expect("4 byte slice"), + ) as usize) + .saturating_sub(4); // subtract the 4 length bytes themselves + + let total = 5 + body_len; // tag(1) + length(4) + body + while self.read_buf.len() < total { + let n = self.stream.read_buf(&mut self.read_buf).await?; + if n == 0 { + bail!("server closed connection while reading message body"); + } + } + + self.read_buf.advance(5); // consume tag + length + let body = self.read_buf.split_to(body_len).freeze(); + Ok((tag, body)) + } + + // ----------------------------------------------------------------------- + // Authentication + // ----------------------------------------------------------------------- + + async fn handle_auth_request( + &mut self, + body: &[u8], + user: &str, + password: &str, + ) -> anyhow::Result<()> { + if body.len() < 4 { + bail!("AuthenticationRequest body too short"); + } + let auth_type = i32::from_be_bytes(body[0..4].try_into().expect("4 bytes")); + + match auth_type { + 0 => { + // AuthenticationOk — nothing to do + } + 3 => { + // AuthenticationCleartextPassword + let mut buf = BytesMut::new(); + frontend::password_message(password.as_bytes(), &mut buf)?; + self.stream.write_all(&buf).await?; + self.stream.flush().await?; + } + 5 => { + // AuthenticationMD5Password — salt is bytes 4..8 + if body.len() < 8 { + bail!("MD5 auth: body too short for salt"); + } + let salt: [u8; 4] = body[4..8].try_into().expect("4 byte salt"); + let hashed = md5_hash(user.as_bytes(), password.as_bytes(), salt); + let mut buf = BytesMut::new(); + frontend::password_message(hashed.as_bytes(), &mut buf)?; + self.stream.write_all(&buf).await?; + self.stream.flush().await?; + } + 10 => { + // AuthenticationSASL (SCRAM-SHA-256) — not implemented here. + // For Docker dev setups, use `md5` or `trust` in pg_hba.conf. + bail!( + "SCRAM-SHA-256 authentication is not supported by the \ + replication module. Configure PostgreSQL to use `md5` \ + or `trust` for this user in pg_hba.conf." + ); + } + other => { + bail!("unsupported authentication type: {other}"); + } + } + + Ok(()) + } +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/// Extract a human-readable message from an `ErrorResponse` body. +/// +/// Fields are null-terminated key-value pairs. We extract the 'M' (message) +/// and 'D' (detail) fields. +fn parse_error_message(body: &[u8]) -> String { + let mut parts = Vec::new(); + let mut i = 0; + while i < body.len() { + let field_type = body[i]; + i += 1; + if field_type == 0 { + break; + } + // Find the null terminator of the field value. + let start = i; + while i < body.len() && body[i] != 0 { + i += 1; + } + let value = String::from_utf8_lossy(&body[start..i]).into_owned(); + i += 1; // skip null terminator + if matches!(field_type, b'M' | b'D' | b'C') { + parts.push(value); + } + } + if parts.is_empty() { + "unknown PostgreSQL error".into() + } else { + parts.join(" — ") + } +} diff --git a/src/pg_replication/wal_parser.rs b/src/pg_replication/wal_parser.rs new file mode 100644 index 0000000..685dba1 --- /dev/null +++ b/src/pg_replication/wal_parser.rs @@ -0,0 +1,505 @@ +use bytes::{Buf, BytesMut}; +use std::collections::HashMap; + +/// PostgreSQL logical replication binary protocol parser (pgoutput, version 1). +/// +/// Parses the raw WAL bytes delivered by `START_REPLICATION … (proto_version '1')`. +/// In protocol version 1 the transaction ID (XID) is **only** present in `Begin` +/// messages; DML messages (Insert / Update / Delete / Truncate) carry no XID. +/// The [`WalEmitter`](super::emitter::WalEmitter) tracks the current XID via the +/// `Begin`/`Commit` boundary. + +// --------------------------------------------------------------------------- +// Schema metadata +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +pub struct ColumnInfo { + /// Bit 0: column is part of the replica identity key. + pub flags: u8, + pub name: String, + pub type_oid: u32, + pub type_modifier: i32, +} + +#[derive(Debug, Clone)] +pub struct RelationInfo { + pub oid: u32, + pub namespace: String, + pub name: String, + /// `d` = default (PK), `n` = nothing, `f` = full row, `i` = index. + pub replica_identity: u8, + pub columns: Vec, +} + +// --------------------------------------------------------------------------- +// Tuple data +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +pub enum ColumnValue { + Null, + UnchangedToasted, + Text(String), + Binary(Vec), +} + +#[derive(Debug, Clone)] +pub struct TupleData { + pub values: Vec, +} + +// --------------------------------------------------------------------------- +// Logical messages (pgoutput v1) +// --------------------------------------------------------------------------- + +#[derive(Debug)] +pub enum LogicalMessage { + /// Transaction start. `xid` identifies the transaction; DML messages that + /// follow do not carry their own XID. + Begin { + /// Final LSN of the transaction (equals the Commit LSN). + final_lsn: u64, + /// Commit timestamp (microseconds since PG epoch 2000-01-01). + timestamp: i64, + /// Transaction ID. + xid: u32, + }, + + /// Schema descriptor sent before the first DML that references a relation + /// in a session (and after schema changes). Stored internally by the parser. + Relation(RelationInfo), + + /// Row inserted into `rel_oid`. + Insert { + rel_oid: u32, + new_tuple: TupleData, + }, + + /// Row updated in `rel_oid`. + /// + /// `key_tuple` is present when REPLICA IDENTITY is the primary key + /// (default). `old_tuple` is present for REPLICA IDENTITY FULL. + Update { + rel_oid: u32, + key_tuple: Option, + old_tuple: Option, + new_tuple: TupleData, + }, + + /// Row deleted from `rel_oid`. + Delete { + rel_oid: u32, + key_tuple: Option, + old_tuple: Option, + }, + + /// Transaction committed successfully. + Commit { + flags: u8, + commit_lsn: u64, + end_lsn: u64, + timestamp: i64, + }, + + /// One or more tables truncated. + Truncate { + num_relations: u32, + options: u8, + rel_oids: Vec, + }, + + /// Custom data-type descriptor (pgoutput emits these before first use). + Type { + type_oid: u32, + namespace: String, + type_name: String, + }, + + /// Application-level logical message emitted via `pg_logical_emit_message`. + Message { + /// Bit 0: message is transactional. + flags: u8, + lsn: u64, + prefix: String, + content: Vec, + }, +} + +// --------------------------------------------------------------------------- +// Parser +// --------------------------------------------------------------------------- + +/// Stateful pgoutput v1 parser. +/// +/// Maintains a cache of `RelationInfo` objects keyed by OID so that +/// Insert/Update/Delete messages can be decoded without extra round-trips. +pub struct MessageParser { + relations: HashMap, + _types: HashMap, +} + +impl MessageParser { + pub fn new() -> Self { + Self { + relations: HashMap::new(), + _types: HashMap::new(), + } + } + + // ----------------------------------------------------------------------- + // Public entry point + // ----------------------------------------------------------------------- + + /// Parse the next logical message from `buf`. + /// + /// `Relation` messages are absorbed into internal state and return `None` + /// so callers only see semantically meaningful messages. + pub fn parse_message(&mut self, buf: &mut BytesMut) -> Result, String> { + if buf.is_empty() { + return Ok(None); + } + + match buf[0] as char { + 'B' => self.parse_begin(buf), + 'C' => self.parse_commit(buf), + 'R' => { + self.parse_relation(buf)?; + Ok(None) + } + 'I' => self.parse_insert(buf), + 'U' => self.parse_update(buf), + 'D' => self.parse_delete(buf), + 'T' => self.parse_truncate(buf), + 'Y' => self.parse_type(buf), + 'M' => self.parse_message_msg(buf), + other => Err(format!( + "unknown pgoutput message type: '{}' (0x{:02x})", + other, buf[0] + )), + } + } + + /// Look up a relation by OID (returns `None` if not yet seen). + pub fn get_relation(&self, oid: u32) -> Option<&RelationInfo> { + self.relations.get(&oid) + } + + // ----------------------------------------------------------------------- + // Low-level readers + // ----------------------------------------------------------------------- + + fn read_u8(buf: &mut BytesMut) -> Result { + if buf.is_empty() { + return Err("unexpected EOF reading u8".into()); + } + Ok(buf.get_u8()) + } + + fn read_i16(buf: &mut BytesMut) -> Result { + if buf.len() < 2 { + return Err("unexpected EOF reading i16".into()); + } + Ok(buf.get_i16()) + } + + fn read_u32(buf: &mut BytesMut) -> Result { + if buf.len() < 4 { + return Err("unexpected EOF reading u32".into()); + } + Ok(buf.get_u32()) + } + + fn read_i32(buf: &mut BytesMut) -> Result { + if buf.len() < 4 { + return Err("unexpected EOF reading i32".into()); + } + Ok(buf.get_i32()) + } + + fn read_u64(buf: &mut BytesMut) -> Result { + if buf.len() < 8 { + return Err("unexpected EOF reading u64".into()); + } + Ok(buf.get_u64()) + } + + fn read_i64(buf: &mut BytesMut) -> Result { + if buf.len() < 8 { + return Err("unexpected EOF reading i64".into()); + } + Ok(buf.get_i64()) + } + + /// Read a null-terminated UTF-8 string. + fn read_string(buf: &mut BytesMut) -> Result { + let mut bytes = Vec::new(); + loop { + if buf.is_empty() { + return Err("unexpected EOF reading null-terminated string".into()); + } + let b = buf.get_u8(); + if b == 0 { + break; + } + bytes.push(b); + } + String::from_utf8(bytes).map_err(|e| format!("invalid UTF-8 in string: {e}")) + } + + // ----------------------------------------------------------------------- + // Message parsers + // ----------------------------------------------------------------------- + + fn parse_begin(&mut self, buf: &mut BytesMut) -> Result, String> { + buf.get_u8(); // 'B' + let final_lsn = Self::read_u64(buf)?; + let timestamp = Self::read_i64(buf)?; + let xid = Self::read_u32(buf)?; + Ok(Some(LogicalMessage::Begin { final_lsn, timestamp, xid })) + } + + fn parse_commit(&mut self, buf: &mut BytesMut) -> Result, String> { + buf.get_u8(); // 'C' + let flags = Self::read_u8(buf)?; + let commit_lsn = Self::read_u64(buf)?; + let end_lsn = Self::read_u64(buf)?; + let timestamp = Self::read_i64(buf)?; + Ok(Some(LogicalMessage::Commit { flags, commit_lsn, end_lsn, timestamp })) + } + + fn parse_relation(&mut self, buf: &mut BytesMut) -> Result<(), String> { + buf.get_u8(); // 'R' + let rel_oid = Self::read_u32(buf)?; + let namespace = Self::read_string(buf)?; + let name = Self::read_string(buf)?; + let replica_identity = Self::read_u8(buf)?; + let num_columns = Self::read_i16(buf)? as usize; + + let mut columns = Vec::with_capacity(num_columns); + for _ in 0..num_columns { + let flags = Self::read_u8(buf)?; + let col_name = Self::read_string(buf)?; + let type_oid = Self::read_u32(buf)?; + let type_modifier = Self::read_i32(buf)?; + columns.push(ColumnInfo { flags, name: col_name, type_oid, type_modifier }); + } + + self.relations.insert( + rel_oid, + RelationInfo { oid: rel_oid, namespace, name, replica_identity, columns }, + ); + Ok(()) + } + + fn parse_insert(&mut self, buf: &mut BytesMut) -> Result, String> { + buf.get_u8(); // 'I' + // pgoutput v1: Relation OID comes directly after type byte (no XID). + let rel_oid = Self::read_u32(buf)?; + + let marker = Self::read_u8(buf)?; + if marker != b'N' { + return Err(format!("INSERT: expected 'N' marker, got '{}'", marker as char)); + } + + let new_tuple = self.parse_tuple_data(buf, rel_oid)?; + Ok(Some(LogicalMessage::Insert { rel_oid, new_tuple })) + } + + fn parse_update(&mut self, buf: &mut BytesMut) -> Result, String> { + buf.get_u8(); // 'U' + // pgoutput v1: no XID in Update. + let rel_oid = Self::read_u32(buf)?; + + let mut key_tuple = None; + let mut old_tuple = None; + + if !buf.is_empty() { + match buf[0] as char { + 'K' => { + buf.get_u8(); + key_tuple = Some(self.parse_tuple_data(buf, rel_oid)?); + } + 'O' => { + buf.get_u8(); + old_tuple = Some(self.parse_tuple_data(buf, rel_oid)?); + } + _ => {} // no optional tuple, proceed to 'N' + } + } + + let marker = Self::read_u8(buf)?; + if marker != b'N' { + return Err(format!("UPDATE: expected 'N' marker, got '{}'", marker as char)); + } + let new_tuple = self.parse_tuple_data(buf, rel_oid)?; + + Ok(Some(LogicalMessage::Update { rel_oid, key_tuple, old_tuple, new_tuple })) + } + + fn parse_delete(&mut self, buf: &mut BytesMut) -> Result, String> { + buf.get_u8(); // 'D' + // pgoutput v1: no XID in Delete. + let rel_oid = Self::read_u32(buf)?; + + let mut key_tuple = None; + let mut old_tuple = None; + + if !buf.is_empty() { + match buf[0] as char { + 'K' => { + buf.get_u8(); + key_tuple = Some(self.parse_tuple_data(buf, rel_oid)?); + } + 'O' => { + buf.get_u8(); + old_tuple = Some(self.parse_tuple_data(buf, rel_oid)?); + } + other => { + return Err(format!("DELETE: expected 'K' or 'O', got '{}'", other)); + } + } + } + + Ok(Some(LogicalMessage::Delete { rel_oid, key_tuple, old_tuple })) + } + + fn parse_truncate(&mut self, buf: &mut BytesMut) -> Result, String> { + buf.get_u8(); // 'T' + // pgoutput v1: no XID in Truncate. + let num_relations = Self::read_u32(buf)?; + let options = Self::read_u8(buf)?; + + let mut rel_oids = Vec::with_capacity(num_relations as usize); + for _ in 0..num_relations { + rel_oids.push(Self::read_u32(buf)?); + } + + Ok(Some(LogicalMessage::Truncate { num_relations, options, rel_oids })) + } + + fn parse_type(&mut self, buf: &mut BytesMut) -> Result, String> { + buf.get_u8(); // 'Y' + // pgoutput v1: no XID in Type. + let type_oid = Self::read_u32(buf)?; + let namespace = Self::read_string(buf)?; + let type_name = Self::read_string(buf)?; + + self._types.insert(type_oid, (namespace.clone(), type_name.clone())); + + Ok(Some(LogicalMessage::Type { type_oid, namespace, type_name })) + } + + fn parse_message_msg(&mut self, buf: &mut BytesMut) -> Result, String> { + buf.get_u8(); // 'M' + // pgoutput v1: no XID in Message. + let flags = Self::read_u8(buf)?; + let lsn = Self::read_u64(buf)?; + let prefix = Self::read_string(buf)?; + let content_len = Self::read_i32(buf)? as usize; + + if buf.len() < content_len { + return Err(format!( + "Message: need {} bytes for content, have {}", + content_len, + buf.len() + )); + } + let content = buf.split_to(content_len).to_vec(); + + Ok(Some(LogicalMessage::Message { flags, lsn, prefix, content })) + } + + // ----------------------------------------------------------------------- + // Tuple decoder + // ----------------------------------------------------------------------- + + fn parse_tuple_data(&self, buf: &mut BytesMut, rel_oid: u32) -> Result { + let num_columns = Self::read_i16(buf)? as usize; + + let relation = self + .relations + .get(&rel_oid) + .ok_or_else(|| format!("unknown relation OID {rel_oid} — missing Relation message?"))?; + + if num_columns != relation.columns.len() { + return Err(format!( + "tuple has {} columns but relation '{}' has {}", + num_columns, + relation.name, + relation.columns.len() + )); + } + + let mut values = Vec::with_capacity(num_columns); + for _ in 0..num_columns { + let marker = Self::read_u8(buf)?; + let value = match marker { + b'n' => ColumnValue::Null, + b'u' => ColumnValue::UnchangedToasted, + b't' => { + let len = Self::read_i32(buf)? as usize; + if buf.len() < len { + return Err(format!( + "text column: need {} bytes, have {}", + len, + buf.len() + )); + } + let data = buf.split_to(len).to_vec(); + let text = String::from_utf8(data) + .map_err(|e| format!("non-UTF-8 text column: {e}"))?; + ColumnValue::Text(text) + } + b'b' => { + let len = Self::read_i32(buf)? as usize; + if buf.len() < len { + return Err(format!( + "binary column: need {} bytes, have {}", + len, + buf.len() + )); + } + ColumnValue::Binary(buf.split_to(len).to_vec()) + } + other => { + return Err(format!( + "unknown column marker '{}'", + other as char + )); + } + }; + values.push(value); + } + + Ok(TupleData { values }) + } +} + +impl Default for MessageParser { + fn default() -> Self { + Self::new() + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parser_starts_empty() { + let p = MessageParser::new(); + assert!(p.relations.is_empty()); + } + + #[test] + fn empty_buf_returns_none() { + let mut p = MessageParser::new(); + let mut buf = BytesMut::new(); + assert!(p.parse_message(&mut buf).unwrap().is_none()); + } +}