livevue-rs/examples/pg_replication/main.rs

//! # 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<Markup> {
    // 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<Message> = 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<livevue_rs::server::SseParams>,
) -> impl IntoResponse {
    livevue_rs::sse_handler(State(state.0), query).await
}

#[derive(Deserialize)]
struct PostMessageBody {
    #[serde(rename = "newMsg")]
    new_msg: Option<String>,
}

async fn post_message_handler(
    State(state): State<(SharedState, RenderState)>,
    axum::extract::Query(params): axum::extract::Query<ActionParams>,
    Json(body): Json<serde_json::Value>,
) -> 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(())
}