db_library/database/

postgres.rs

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sqlx::{postgres::PgListener, Executor, PgPool};
use std::sync::Arc;
use tokio::{
    sync::{
        mpsc::{self, Receiver, Sender},
        Mutex,
    },
    task::JoinHandle,
};
use tracing::{error, info};

use crate::config::DBListenerError;
use crate::utils::stringify;

use super::{DBListenerTrait, EventType};

#[derive(Debug, Clone)]
pub struct PostgresTableListener {
    pub pool: Arc<PgPool>,
    pub table_name: String,
    pub columns: Vec<String>,
    pub sender: Sender<Value>,
    pub receiver: Arc<Mutex<tokio::sync::mpsc::Receiver<Value>>>,
    pub table_identifier: String,
    pub pg_trigger_name: String,
    pub pg_function_name: String,
    pub pg_column_updates_name: String,
    pub events: Vec<EventType>,
}

impl PostgresTableListener {
    pub async fn new(
        pool: Arc<PgPool>,
        table_name: &str,
        columns: Vec<String>,
        table_identifier: &str,
        events: Vec<EventType>,
    ) -> Result<Self, DBListenerError> {
        let uniquekey_uuid = uuid::Uuid::new_v4().to_string();
        let uniquekey = uniquekey_uuid.replace("-", "_");

        let pg_trigger_name = format!("{}_{}_trigger", table_name, &uniquekey);
        let pg_function_name = format!("{}_{}_function", table_name, &uniquekey);
        let pg_column_updates_name = format!("{}_{}_column_updates", table_name, &uniquekey);

        let (sender, receiver) = mpsc::channel::<Value>(100); // Channel with buffer size 100

        let postgres_table_listener = Self {
            pool,
            table_name: table_name.to_string(),
            columns: columns.into_iter().map(|c| c.to_string()).collect(),
            table_identifier: table_identifier.to_string(),
            sender,
            receiver: Arc::new(Mutex::new(receiver)),
            pg_trigger_name,
            pg_function_name,
            pg_column_updates_name,
            events,
        };

        postgres_table_listener.verify_members().await?;

        Ok(postgres_table_listener)
    }

    async fn verify_members(&self) -> Result<(), DBListenerError> {
        info!("--> Verifying table and columns");
        // Verify that the table exists
        let table_exists = sqlx::query_as::<_, (bool,)>(
            r#"
            SELECT EXISTS (
                SELECT 1
                FROM information_schema.tables
                WHERE table_name = $1
            );
            "#,
        )
        .bind(&self.table_name)
        .fetch_one(&*self.pool)
        .await
        .map_err(|e| {
            DBListenerError::ListenerVerifyError(format!(
                "Failed to verify table existence : {:?}",
                e
            ))
        })?;

        info!("table exists : {:#?}", table_exists);

        if !table_exists.0 {
            return Err(DBListenerError::ListenerVerifyError(format!(
                "Table '{}' does not exist",
                &self.table_name
            )));
        }

        // verify all columns exist
        for column in &self.columns {
            let column_exists = sqlx::query_as::<_, (bool,)>(
                r#"
                SELECT EXISTS (
                    SELECT 1
                    FROM information_schema.columns
                    WHERE table_name = $1
                    AND column_name = $2
                );
                "#,
            )
            .bind(&self.table_name)
            .bind(column)
            .fetch_one(&*self.pool)
            .await
            .map_err(|e| {
                DBListenerError::ListenerVerifyError(format!(
                    "Failed to verify column existence '{}': {:?}",
                    column, e
                ))
            })?;

            if !column_exists.0 {
                return Err(DBListenerError::ListenerVerifyError(format!(
                    "Column '{}' does not exist",
                    column
                )));
            }
        }

        // Verify that the table identifier exists
        let table_identifier_exists = sqlx::query_as::<_, (bool,)>(
            r#"
            SELECT EXISTS (
                SELECT 1
                FROM information_schema.columns
                WHERE table_name = $1
                AND column_name = $2
            );
            "#,
        )
        .bind(&self.table_name)
        .bind(&self.table_identifier)
        .fetch_one(&*self.pool)
        .await
        .map_err(|e| {
            DBListenerError::ListenerVerifyError(format!(
                "Failed to verify table identifier '{}': {:?}",
                &self.table_identifier, e
            ))
        })?;

        if !table_identifier_exists.0 {
            return Err(DBListenerError::ListenerVerifyError(format!(
                "Table identifier '{}' does not exist",
                &self.table_identifier
            )));
        }

        info!("✅ Table and columns verified successfully");
        Ok(())
    }

    async fn create_trigger_and_function(&self) -> Result<(), DBListenerError> {
        let table_identifier = &self.table_identifier;

        // Create the function with clean INSERT handling
        let create_function = format!(
            r#"
                CREATE OR REPLACE FUNCTION {function_name}()
                RETURNS TRIGGER AS $$
                BEGIN
                    IF TG_OP = 'INSERT' THEN 
                        {insert_blocks}
                    ELSIF TG_OP = 'UPDATE' THEN 
                        {update_blocks}
                    ELSIF TG_OP = 'DELETE' THEN
                        {delete_blocks}
                    END IF;
                    RETURN NEW;
                END;
                $$ LANGUAGE plpgsql;
            "#,
            function_name = self.pg_function_name,
            insert_blocks = self
                .columns
                .iter()
                .map(|col| {
                    format!(
                        r#"
                        IF NEW.{col} IS NOT NULL THEN
                            PERFORM pg_notify(
                                '{column_updates}',
                                json_build_object(
                                    'operation', TG_OP,
                                    'table', TG_TABLE_NAME,
                                    'column', '{col}',
                                    'id', NEW.{table_identifier},
                                    'new_value', NEW.{col},
                                    'timestamp', NOW()
                                )::text
                            );
                        END IF;
                    "#,
                        column_updates = self.pg_column_updates_name
                    )
                })
                .collect::<Vec<_>>()
                .join("\n"),
            update_blocks = self
                .columns
                .iter()
                .map(|col| {
                    format!(
                        r#"
                        IF NEW.{col} IS DISTINCT FROM OLD.{col} THEN
                            PERFORM pg_notify(
                                '{column_updates}',
                                json_build_object(
                                    'operation', TG_OP,
                                    'table', TG_TABLE_NAME,
                                    'column', '{col}',
                                    'id', NEW.{table_identifier},
                                    'old_value', OLD.{col},
                                    'new_value', NEW.{col},
                                    'timestamp', NOW()
                                )::text
                            );
                        END IF;
                    "#,
                        column_updates = self.pg_column_updates_name,
                    )
                })
                .collect::<Vec<_>>()
                .join("\n"),
            delete_blocks = self
                .columns
                .iter()
                .map(|col| {
                    format!(
                        r#"
                        IF OLD.{col} IS NOT NULL THEN
                            PERFORM pg_notify(
                                '{column_updates}',
                                json_build_object(
                                    'operation', TG_OP,
                                    'table', TG_TABLE_NAME,
                                    'column', '{col}',
                                    'id', OLD.{table_identifier},
                                    'old_value', OLD.{col},
                                    'timestamp', NOW()
                                )::text
                            );
                        END IF;
                    "#,
                        column_updates = self.pg_column_updates_name,
                    )
                })
                .collect::<Vec<_>>()
                .join("\n"),
        );

        // Execute function creation
        self.pool
            .execute(create_function.as_str())
            .await
            .map_err(|e| {
                DBListenerError::CreationError(format!(
                    "Failed to execute function creation : {:#?}",
                    e
                ))
            })?;

        let events_list = EventType::get_events_list(&self.events, &self.columns.join(", "));

        // Create single trigger for INSERT
        let create_trigger = format!(
            r#"
            DO $$ 
            BEGIN 
                IF NOT EXISTS (
                    SELECT 1 
                    FROM pg_trigger 
                    WHERE tgname = '{trigger_name}'
                ) THEN 
                    CREATE TRIGGER {trigger_name}
                    AFTER {events_list} ON {table_name}
                    FOR EACH ROW
                    EXECUTE FUNCTION {function_name}();
                END IF;
            END $$;
            "#,
            trigger_name = self.pg_trigger_name,
            table_name = self.table_name,
            function_name = self.pg_function_name
        );

        // Execute trigger creation
        self.pool
            .execute(create_trigger.as_str())
            .await
            .map_err(|e| {
                DBListenerError::CreationError(format!(
                    "Failed to execute trigger creation : {:#?}",
                    e
                ))
            })?;

        Ok(())
    }

    async fn drop_trigger_and_function(&self) -> Result<(), DBListenerError> {
        let trigger_name = &self.pg_trigger_name;
        let function_name = &self.pg_function_name;
        let table_name = &self.table_name;

        // Drop the trigger if it exists
        let drop_trigger = format!(
            r#"
            DO $$
            BEGIN
                IF EXISTS (
                    SELECT 1
                    FROM pg_trigger
                    WHERE tgname = '{trigger_name}'
                ) THEN
                    DROP TRIGGER {trigger_name} ON {table_name};
                END IF;
            END $$;
            "#,
            trigger_name = trigger_name,
            table_name = table_name
        );

        self.pool
            .execute(drop_trigger.as_str())
            .await
            .map_err(|e| {
                DBListenerError::DeletionError(format!(
                    "Failed to execute trigger deletion : {:#?}",
                    e
                ))
            })?;

        // Drop the function if it exists
        let drop_function = format!(
            r#"
            DO $$
            BEGIN
                IF EXISTS (
                    SELECT 1
                    FROM pg_proc
                    WHERE proname = '{function_name}'
                ) THEN
                    DROP FUNCTION {function_name}();
                END IF;
            END $$;
            "#,
            function_name = function_name
        );

        self.pool
            .execute(drop_function.as_str())
            .await
            .map_err(|e| {
                DBListenerError::DeletionError(format!(
                    "Failed to execute trigger deletion : {:#?}",
                    e
                ))
            })?;

        info!("Trigger and function removed for table: {}", table_name);
        Ok(())
    }
}

#[async_trait]
impl DBListenerTrait for PostgresTableListener {
    async fn start(
        &self,
    ) -> Result<(Arc<Mutex<Receiver<Value>>>, JoinHandle<()>), DBListenerError> {
        // Create trigger and function if they don’t exist
        self.create_trigger_and_function()
            .await
            .map_err(|e| DBListenerError::ListenerError(e.to_string()))?;

        let listener = PgListener::connect_with(&self.pool)
            .await
            .map_err(|e| DBListenerError::ListenerError(e.to_string()))?;

        let listener = Arc::new(Mutex::new(listener));

        {
            let mut locked_listener = listener.lock().await;
            locked_listener
                .listen(&self.pg_column_updates_name)
                .await
                .map_err(|e| DBListenerError::ListenerError(e.to_string()))?;
        }

        info!(
            "Listening for column update notifications on {}",
            self.table_name
        );

        let listener_clone = Arc::clone(&listener);
        let sender_clone = self.sender.clone();

        let handle = tokio::spawn(async move {
            info!("Listener spawned and waiting for notifications");

            loop {
                let mut locked_listener = listener_clone.lock().await;
                match locked_listener.recv().await {
                    Ok(notification) => {
                        if let Ok(payload) = serde_json::from_str::<Value>(&notification.payload())
                        {
                            if let Err(e) = sender_clone.send(payload).await {
                                error!("Failed to send payload to channel: {:?}", e);
                            }
                        }
                    }
                    Err(e) => {
                        error!("Listener encountered an error: {:?}", e);
                        break;
                    }
                }
            }
        });

        Ok((Arc::clone(&self.receiver), handle))
    }

    async fn stop(&self) -> Result<(), DBListenerError> {
        self.drop_trigger_and_function()
            .await
            .map_err(|e| DBListenerError::ListenerError(e.to_string()))?;
        Ok(())
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct PgNotify {
    pub operation: Option<String>,
    pub table: String,

    #[serde(deserialize_with = "stringify")]
    pub id: String, // Convert integer to string during deserialization

    pub old_value: Option<String>,
    pub new_value: Option<String>,
    pub timestamp: String,
}