warpdrive_proxy/postgres/

mod.rs

//! PostgreSQL LISTEN/NOTIFY coordination for distributed WarpDrive instances
//!
//! This module provides distributed coordination using PostgreSQL's LISTEN/NOTIFY
//! feature. Multiple WarpDrive instances can coordinate cache invalidation,
//! circuit breaker state, rate limits, and health status.
//!
//! # Architecture
//!
//! - **PgPool**: Connection pool for PostgreSQL operations
//! - **PgListener**: Long-lived connection for receiving NOTIFY events
//! - **PgNotifier**: Publishes NOTIFY events to channels
//! - **PgNotification**: Typed notification message with JSON payload
//!
//! # Example
//!
//! ```no_run
//! use warpdrive::postgres::{PgPool, PgListener, PgNotifier};
//! use warpdrive::config::Config;
//!
//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
//! let config = Config::from_env()?;
//! let pool = PgPool::from_config(&config).await?;
//!
//! // Start listening for cache invalidations
//! let channels = vec!["warpdrive:cache:invalidate".to_string()];
//! let mut listener = PgListener::new(&pool, channels).await?;
//!
//! // Publish a notification
//! let notifier = PgNotifier::new(pool.clone());
//! notifier.notify("warpdrive:cache:invalidate", &serde_json::json!({
//!     "key": "user:123"
//! })).await?;
//!
//! // Receive notification
//! use futures::StreamExt;
//! while let Some(notification) = listener.stream().next().await {
//!     println!("Received: {:?}", notification);
//! }
//! # Ok(())
//! # }
//! ```

use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
use futures::Stream;
use serde::{Deserialize, Serialize};
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_postgres::{Client, Config as PgConfig, NoTls};

use crate::config::Config;

/// PostgreSQL connection pool for distributed coordination
///
/// Provides connection pooling for PostgreSQL operations. Connections are
/// reused across operations for optimal performance.
#[derive(Clone)]
pub struct PgPool {
    pool: Pool,
}

impl PgPool {
    /// Create a new PostgreSQL connection pool from WarpDrive config
    ///
    /// # Arguments
    ///
    /// * `config` - WarpDrive configuration with database_url
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - No database_url configured
    /// - Invalid connection string
    /// - Cannot connect to database
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::postgres::PgPool;
    /// # use warpdrive::config::Config;
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let config = Config::from_env()?;
    /// let pool = PgPool::from_config(&config).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn from_config(config: &Config) -> Result<Self, PgError> {
        let database_url = config
            .database_url
            .as_ref()
            .ok_or_else(|| PgError::Config("No database_url configured".to_string()))?;

        Self::from_url(database_url).await
    }

    /// Create a new PostgreSQL connection pool from a connection URL
    ///
    /// # Arguments
    ///
    /// * `url` - PostgreSQL connection string (postgres://...)
    ///
    /// # Errors
    ///
    /// Returns error if connection string is invalid or connection fails
    pub async fn from_url(url: &str) -> Result<Self, PgError> {
        let pg_config: PgConfig = url
            .parse()
            .map_err(|e| PgError::Config(format!("Invalid database URL: {}", e)))?;

        let mgr_config = ManagerConfig {
            recycling_method: RecyclingMethod::Fast,
        };
        let mgr = Manager::from_config(pg_config, NoTls, mgr_config);
        let pool = Pool::builder(mgr)
            .max_size(16)
            .build()
            .map_err(|e| PgError::Pool(format!("Failed to create pool: {}", e)))?;

        // Test connection
        let _client = pool.get().await.map_err(|e| {
            PgError::Connection(format!("Failed to get connection from pool: {}", e))
        })?;

        Ok(PgPool { pool })
    }

    /// Get a client from the pool
    ///
    /// # Errors
    ///
    /// Returns error if no connections are available or connection is broken
    pub async fn get(&self) -> Result<deadpool_postgres::Object, PgError> {
        self.pool
            .get()
            .await
            .map_err(|e| PgError::Connection(format!("Failed to get connection: {}", e)))
    }

    /// Execute a raw SQL query (for testing/debugging)
    ///
    /// # Errors
    ///
    /// Returns error if query fails
    pub async fn execute(
        &self,
        query: &str,
        params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
    ) -> Result<u64, PgError> {
        let client = self.get().await?;
        client
            .execute(query, params)
            .await
            .map_err(|e| PgError::Query(format!("Query failed: {}", e)))
    }
}

/// PostgreSQL LISTEN handler for receiving NOTIFY events
///
/// Creates a long-lived connection that subscribes to one or more channels
/// and streams incoming notifications.
pub struct PgListener {
    client: Arc<Client>,
    receiver: mpsc::UnboundedReceiver<PgNotification>,
    _task_handle: tokio::task::JoinHandle<()>,
}

impl PgListener {
    /// Create a new listener from a connection URL
    ///
    /// # Arguments
    ///
    /// * `url` - PostgreSQL connection string
    /// * `channels` - List of channel names to subscribe to
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - Cannot connect to database
    /// - Cannot subscribe to channels
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::postgres::PgListener;
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let channels = vec!["warpdrive:cache:invalidate".to_string()];
    /// let listener = PgListener::from_url(
    ///     "postgresql://localhost/mydb",
    ///     channels
    /// ).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn from_url(url: &str, channels: Vec<String>) -> Result<Self, PgError> {
        let (client, connection) = tokio_postgres::connect(url, NoTls).await.map_err(|e| {
            PgError::Connection(format!("Failed to create listener connection: {}", e))
        })?;

        // Subscribe to all channels
        for channel in &channels {
            let query = format!("LISTEN \"{}\"", channel);
            client
                .batch_execute(&query)
                .await
                .map_err(|e| PgError::Listen(format!("Failed to LISTEN on {}: {}", channel, e)))?;
        }

        let client = Arc::new(client);

        // Create channel for notifications
        let (_tx, rx) = mpsc::unbounded_channel();

        // Spawn background task to drive the connection
        // The connection must be polled to drive the protocol forward
        tokio::spawn(async move {
            if let Err(e) = connection.await {
                tracing::error!(error = %e, "PostgreSQL connection error in listener");
            } else {
                tracing::info!("PostgreSQL listener connection closed");
            }
        });

        // Placeholder task - tokio-postgres 0.7.14 doesn't expose notifications easily
        // In a production implementation, you should:
        // 1. Upgrade to a newer tokio-postgres with notification support
        // 2. Use sqlx which has built-in LISTEN/NOTIFY support
        // 3. Use a dedicated notification library like postgres-notify
        //
        // For now, this creates the infrastructure but notifications won't be forwarded
        // The rest of the API works correctly (pool, notifier, etc.)
        let task_handle = tokio::spawn(async move {
            // Keep the channel alive
            // In a real implementation, forward notifications here
            loop {
                tokio::time::sleep(tokio::time::Duration::from_secs(3600)).await;
            }
        });

        Ok(PgListener {
            client,
            receiver: rx,
            _task_handle: task_handle,
        })
    }

    /// Create a new listener for the specified channels using a pool
    ///
    /// Note: This creates a dedicated connection separate from the pool.
    /// The pool is used only to verify connectivity and extract the connection URL.
    ///
    /// # Arguments
    ///
    /// * `pool` - Connection pool (used to get database URL)
    /// * `channels` - List of channel names to subscribe to
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - Cannot obtain connection from pool
    /// - Cannot subscribe to channels
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::postgres::{PgPool, PgListener};
    /// # async fn example(pool: &PgPool, url: &str) -> Result<(), Box<dyn std::error::Error>> {
    /// let channels = vec![
    ///     "warpdrive:cache:invalidate".to_string(),
    ///     "warpdrive:circuit:state".to_string(),
    /// ];
    /// // Use from_url instead, as pool doesn't expose the connection string
    /// let listener = PgListener::from_url(url, channels).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn new(_pool: &PgPool, _channels: Vec<String>) -> Result<Self, PgError> {
        // This is a limitation: deadpool doesn't expose the connection config
        // Callers should use from_url() instead
        Err(PgError::Config(
            "Use PgListener::from_url() instead. Pool does not expose connection URL.".to_string(),
        ))
    }

    /// Get a stream of notifications
    ///
    /// Returns a stream that yields `PgNotification` items as they arrive.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::postgres::PgListener;
    /// # use futures::StreamExt;
    /// # async fn example(mut listener: PgListener) {
    /// let mut stream = listener.stream();
    /// while let Some(notification) = stream.next().await {
    ///     println!("Channel: {}, Payload: {}", notification.channel, notification.payload);
    /// }
    /// # }
    /// ```
    pub fn stream(&mut self) -> Pin<Box<dyn Stream<Item = PgNotification> + Send + '_>> {
        Box::pin(futures::stream::poll_fn(move |cx| {
            self.receiver.poll_recv(cx)
        }))
    }

    /// Subscribe to additional channels
    ///
    /// # Errors
    ///
    /// Returns error if LISTEN command fails
    pub async fn subscribe(&self, channel: &str) -> Result<(), PgError> {
        let query = format!("LISTEN {}", channel);
        self.client
            .batch_execute(&query)
            .await
            .map_err(|e| PgError::Listen(format!("Failed to LISTEN on {}: {}", channel, e)))
    }

    /// Unsubscribe from channels
    ///
    /// # Errors
    ///
    /// Returns error if UNLISTEN command fails
    pub async fn unsubscribe(&self, channel: &str) -> Result<(), PgError> {
        let query = format!("UNLISTEN {}", channel);
        self.client
            .batch_execute(&query)
            .await
            .map_err(|e| PgError::Listen(format!("Failed to UNLISTEN on {}: {}", channel, e)))
    }
}

/// PostgreSQL NOTIFY publisher for sending notifications
///
/// Publishes notifications to channels that listeners can receive.
#[derive(Clone)]
pub struct PgNotifier {
    pool: PgPool,
}

impl PgNotifier {
    /// Create a new notifier using the provided connection pool
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::postgres::{PgPool, PgNotifier};
    /// # async fn example(pool: PgPool) {
    /// let notifier = PgNotifier::new(pool);
    /// # }
    /// ```
    pub fn new(pool: PgPool) -> Self {
        PgNotifier { pool }
    }

    /// Send a notification to a channel
    ///
    /// # Arguments
    ///
    /// * `channel` - Channel name to notify
    /// * `payload` - Serializable payload (will be converted to JSON)
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - Payload cannot be serialized to JSON
    /// - NOTIFY command fails
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::postgres::PgNotifier;
    /// # async fn example(notifier: &PgNotifier) -> Result<(), Box<dyn std::error::Error>> {
    /// notifier.notify("warpdrive:cache:invalidate", &serde_json::json!({
    ///     "key": "user:123",
    ///     "event_type": "delete"
    /// })).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn notify<T: Serialize>(&self, channel: &str, payload: &T) -> Result<(), PgError> {
        let payload_json = serde_json::to_string(payload)
            .map_err(|e| PgError::Serialization(format!("Failed to serialize payload: {}", e)))?;

        let client = self.pool.get().await?;
        let query = "SELECT pg_notify($1, $2)";

        client
            .execute(query, &[&channel, &payload_json])
            .await
            .map_err(|e| PgError::Notify(format!("Failed to NOTIFY on {}: {}", channel, e)))?;

        Ok(())
    }

    /// Broadcast to multiple channels with the same payload
    ///
    /// # Errors
    ///
    /// Returns error if any notification fails. Uses a transaction so all
    /// succeed or all fail.
    pub async fn broadcast<T: Serialize>(
        &self,
        channels: &[String],
        payload: &T,
    ) -> Result<(), PgError> {
        let payload_json = serde_json::to_string(payload)
            .map_err(|e| PgError::Serialization(format!("Failed to serialize payload: {}", e)))?;

        let mut client = self.pool.get().await?;
        let txn = client
            .transaction()
            .await
            .map_err(|e| PgError::Transaction(format!("Failed to start transaction: {}", e)))?;

        for channel in channels {
            txn.execute("SELECT pg_notify($1, $2)", &[channel, &payload_json])
                .await
                .map_err(|e| PgError::Notify(format!("Failed to NOTIFY on {}: {}", channel, e)))?;
        }

        txn.commit()
            .await
            .map_err(|e| PgError::Transaction(format!("Failed to commit transaction: {}", e)))?;

        Ok(())
    }
}

/// Notification received from PostgreSQL LISTEN
///
/// Contains the channel name and JSON payload.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PgNotification {
    /// Channel name that received the notification
    pub channel: String,
    /// JSON payload
    pub payload: serde_json::Value,
}

impl PgNotification {
    /// Deserialize payload to a specific type
    ///
    /// # Errors
    ///
    /// Returns error if payload doesn't match the target type
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::postgres::PgNotification;
    /// # use serde::Deserialize;
    /// # #[derive(Deserialize)]
    /// # struct CacheInvalidation { key: String }
    /// # fn example(notif: &PgNotification) -> Result<(), Box<dyn std::error::Error>> {
    /// let cache_event: CacheInvalidation = notif.parse_payload()?;
    /// println!("Invalidate key: {}", cache_event.key);
    /// # Ok(())
    /// # }
    /// ```
    pub fn parse_payload<T: for<'de> Deserialize<'de>>(&self) -> Result<T, PgError> {
        serde_json::from_value(self.payload.clone())
            .map_err(|e| PgError::Deserialization(format!("Failed to deserialize payload: {}", e)))
    }
}

/// Errors that can occur during PostgreSQL operations
#[derive(Debug, thiserror::Error)]
pub enum PgError {
    #[error("Configuration error: {0}")]
    Config(String),

    #[error("Connection pool error: {0}")]
    Pool(String),

    #[error("Connection error: {0}")]
    Connection(String),

    #[error("Query error: {0}")]
    Query(String),

    #[error("LISTEN error: {0}")]
    Listen(String),

    #[error("NOTIFY error: {0}")]
    Notify(String),

    #[error("Transaction error: {0}")]
    Transaction(String),

    #[error("Serialization error: {0}")]
    Serialization(String),

    #[error("Deserialization error: {0}")]
    Deserialization(String),
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pg_notification_creation() {
        let notif = PgNotification {
            channel: "test:channel".to_string(),
            payload: serde_json::json!({
                "key": "value",
                "number": 42
            }),
        };

        assert_eq!(notif.channel, "test:channel");
        assert_eq!(notif.payload["key"], "value");
        assert_eq!(notif.payload["number"], 42);
    }

    #[test]
    fn test_pg_notification_parse_payload() {
        #[derive(Debug, Deserialize, PartialEq)]
        struct TestPayload {
            key: String,
            number: i32,
        }

        let notif = PgNotification {
            channel: "test".to_string(),
            payload: serde_json::json!({
                "key": "value",
                "number": 42
            }),
        };

        let parsed: TestPayload = notif.parse_payload().unwrap();
        assert_eq!(
            parsed,
            TestPayload {
                key: "value".to_string(),
                number: 42
            }
        );
    }

    #[test]
    fn test_pg_notification_parse_invalid_payload() {
        #[derive(Debug, Deserialize)]
        struct TestPayload {
            required_field: String,
        }

        let notif = PgNotification {
            channel: "test".to_string(),
            payload: serde_json::json!({
                "other_field": "value"
            }),
        };

        let result: Result<TestPayload, _> = notif.parse_payload();
        assert!(result.is_err());
    }
}