degen_sql/db/postgres/

postgres_db.rs

use tokio_postgres::Client;
use crate::db::postgres::models::model::PostgresModelError;
use tokio::time::timeout;
use tokio::time::Duration;
use tokio::time::sleep;
use log::info;
use tokio;
use tokio_postgres::{Error as PostgresError, NoTls};

use std::error::Error;
use tokio_postgres_migration::Migration;

use dotenvy::dotenv;
use std::env;

use std::fs;
use std::str;

type MigrationDefinition = (String, String);

#[derive(Clone)]
struct Migrations {
    up: Vec<MigrationDefinition>,
    down: Vec<MigrationDefinition>,
}

pub trait MigrationAsStr {
    fn to_str(&self) -> (&str, &str);
}

impl MigrationAsStr for MigrationDefinition {
    fn to_str(&self) -> (&str, &str) {
        return (self.0.as_str(), self.1.as_str());
    }
}

pub struct Database {

     pool: deadpool_postgres::Pool, // Or other pool implementation

  //  pub client: Option<  tokio_postgres::Client > ,
    pub migrations_dir_path: Option<String>,
    pub connection_url:  String  , 
   
    pub max_reconnect_attempts: u32, 
    pub timeout_duration: Duration, 
}






#[derive(Debug)]
pub enum DatabaseError {
     ConnectionFailed ,
    PoolCreationFailed(String),
    QueryFailed(tokio_postgres::Error),

    RowParseError(String) , 

    PostgresError(tokio_postgres::Error),
    PoolError(deadpool::managed::PoolError<tokio_postgres::Error>),
}

impl From<tokio_postgres::Error> for DatabaseError {
    fn from(error: tokio_postgres::Error) -> Self {
        DatabaseError::PostgresError(error)
    }
}

impl From<deadpool::managed::PoolError<tokio_postgres::Error>> for DatabaseError {
    fn from(error: deadpool::managed::PoolError<tokio_postgres::Error>) -> Self {
        DatabaseError::PoolError(error)
    }
}

// First implement Display for your error type
impl std::fmt::Display for DatabaseError {
    fn fmt(&self, f: &mut  std::fmt::Formatter<'_>) ->  std::fmt::Result {
        match self {
            DatabaseError::ConnectionFailed  => write!(f, "Database connection failed"),
            DatabaseError::PoolCreationFailed(msg) => write!(f, "Connection pool creation failed: {}", msg),
            DatabaseError::QueryFailed(err) => write!(f, "Database query failed: {}", err),
            DatabaseError::PoolError(err) => write!(f, "Pool error: {}", err),
            DatabaseError::PostgresError( err ) =>  write!(f, "Postgres error: {}", err),
            DatabaseError::RowParseError( msg ) => write!(f, "row parse error: {}", msg),
            // Handle other variants
        }
    }
}

// Then implement the Error trait
impl Error for DatabaseError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            DatabaseError::QueryFailed(err) => Some(err),
            // For PoolError, you'd need to check if it implements Error
            // Otherwise return None for this variant
            _ => None,
        }
    }
}




#[derive(Debug,Clone)]
pub struct DatabaseCredentials {
    pub db_name: String,
    pub db_host: String,
    pub db_user: String,
    pub db_password: String,
}

impl Default for DatabaseCredentials {
    fn default() -> Self {
        Self {
            db_name: "postgres".into(),
            db_host: "localhost".into(),
            db_user: "postgres".into(),
            db_password: "postgres".into(),
        }
    }
}

impl DatabaseCredentials {
    pub fn from_env() -> Self {
        //dotenv().expect(".env file not found"); //need to run this beforehand !! 

        Self {
            db_name: env::var("DB_NAME").unwrap_or("postgres".into()),
            db_host: env::var("DB_HOST").unwrap_or("localhost".into()),
            db_user: env::var("DB_USER").unwrap_or("postgres".into()),
            db_password: env::var("DB_PASSWORD").unwrap_or("postgres".into()),
        }
    }

    pub fn build_connection_url(&self) -> String {
        return format!(
            "postgres://{}:{}@{}/{}",
            self.db_user, self.db_password, self.db_host, self.db_name
        );
    }
}

enum PostgresInputType {
    Query,
    QueryOne,
    Execute,
}

struct PostgresInput<'a> {
    input_type: PostgresInputType,
    query: String,
    params: &'a [&'a (dyn tokio_postgres::types::ToSql + Sync)],
}

impl Database {
    pub fn new(
        conn_url: String, 
        migrations_dir_path: Option<String>,
    ) -> Result<Database, DatabaseError> {
        // Parse the connection config from the URL
        let config: tokio_postgres::Config = conn_url.parse()
            .map_err(|_e| DatabaseError::ConnectionFailed )?;
            
        // Create a manager using the config
        let manager = deadpool_postgres::Manager::new(config, tokio_postgres::NoTls);
        
        // Create the pool with builder pattern
        let pool = deadpool_postgres::Pool::builder(manager)
            .max_size(16)
            .build()
            .map_err(|e| DatabaseError::PoolCreationFailed(e.to_string()))?;
        
        Ok(Database {
            pool,
            migrations_dir_path,
            connection_url: conn_url,
            max_reconnect_attempts: 3,
            timeout_duration: Duration::from_secs(5)
        })
    }
}


impl Database {

   

    pub async fn connect(
       // credentials: DatabaseCredentials,
       &  self 
    ) -> Result<Client, PostgresError> {
        // Define the connection URL.
       // let conn_url = credentials.build_connection_url();

        info!("Connecting to db: {}", self.connection_url);

        let (client, connection) = tokio_postgres::connect(&self.connection_url, NoTls).await?;

        // The connection object performs the actual communication with the database,
        // so spawn it off to run on its own.
        tokio::spawn(async move {
            //this is a blocking call i think !!!
            if let Err(e) = connection.await {
                eprintln!("postgres connection error: {}", e);
            }
        });

     //   self.client = Some(client);

        Ok( client )
    }


/*
    pub async fn reconnect(&mut self  ) -> Result<Option< Client >, PostgresError> {
        let max_retries = 5;
        let mut attempt = 0;
         let conn_url = self.connection_url.clone() ;

        while attempt < max_retries {
            info!("Attempt {}: Reconnecting to database...", attempt + 1);

            match tokio_postgres::connect(&conn_url, NoTls).await {
                Ok((client, connection)) => {
                    // Spawn a task to keep the connection alive
                    tokio::spawn(async move {
                        if let Err(e) = connection.await {
                            eprintln!("postgres connection error: {}", e);
                        }
                    });

                //    self.client = Some(client); // Replace old client with the new one
                    info!("Reconnection successful.");
                    return Ok( Some(client) );
                }
                Err(e) => {
                  
                    attempt += 1;

                    if attempt == max_retries {
                        return Err( e.into() )
                    }
                      eprintln!("Reconnection failed: {}. Retrying...", e);
                    sleep(Duration::from_secs(2_u64.pow(attempt))).await; // Exponential backoff
                }
            }
        }

       Ok(None)  //should error ? 
    }*/


    fn read_migration_files(migrations_dir_path: Option<String>) -> Migrations {
        let mut migrations = Migrations {
            up: Vec::new(),
            down: Vec::new(),
        };

        let migrations_dir =
            migrations_dir_path.unwrap_or("./src/db/postgres/migrations".to_string());
        let migration_dir_files = fs::read_dir(&migrations_dir).expect("Failed to read directory");

        for file in migration_dir_files {
            let file = file.expect("Failed to read migration file");

            let path = file.path();
            let filename = path.file_stem().unwrap().to_str().unwrap();

            let filename_without_extension: &str = filename.split('.').next().unwrap();

            // Read file contents
            let contents = fs::read_to_string(file.path()).expect("Failed to read file contents");

            //let contents = str::from_utf8(file.contents()).unwrap();

            info!("File name: {}", filename);

            if filename.contains(".down") {
                info!("File contents: {}", contents);
                migrations
                    .down
                    .push((filename_without_extension.into(), contents.clone()));
            }

            if filename.contains(".up") {
                info!("File contents: {}", contents);
                migrations
                    .up
                    .push((filename_without_extension.into(), contents.clone()));
            }
        }
        
            
        // Sort `up` migrations in ascending alphabetical order
        migrations.up.sort_by(|a, b| a.0.cmp(&b.0));
        
        // Sort `down` migrations in descending alphabetical order
        migrations.down.sort_by(|a, b| b.0.cmp(&a.0));
                

        return migrations;
    }

    pub async fn migrate(&mut self) -> Result<(), Box<dyn Error>> {
        let client = &mut self.connect().await?;

        let migrations_dir_path = self.migrations_dir_path.clone();
        let mut migrations: Migrations = Self::read_migration_files(migrations_dir_path);

        for up_migration in migrations.up.drain(..) {
            println!("migrating {} {} ", up_migration.0, up_migration.1);
            let migration = Migration::new("migrations".to_string());

            // execute non existing migrations
            migration.up(client, &[up_migration.to_str()]).await?;
        }

        // ...
        Ok(())
    }

    //basically need to do the DOWN migrations and also delete some records from the migrations table
    //need to read from the migrations table with MigrationsModel::find
    pub async fn rollback(&mut self) -> Result<(), Box<dyn Error>> {
        Ok(())
    }

    pub async fn rollback_full(&mut self) -> Result<(), Box<dyn Error>> {
        let migrations_dir_path = self.migrations_dir_path.clone();

        let mut migrations: Migrations = Self::read_migration_files(migrations_dir_path);

            let client = &mut self.connect().await?;

        for down_migration in migrations.down.drain(..)
         {
            println!("migrating {}", down_migration.0);
            let migration = Migration::new("migrations".to_string());
            // execute non existing migrations
            migration.down(client, &[down_migration.to_str()]).await?;
        }

        Ok(())
    }

     pub async fn query(
        &self,
        query: &str,
        params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
    ) -> Result<Vec<tokio_postgres::Row>, DatabaseError> {

        /*let client = &self.connect().await?;

        let rows = client.query(query, params).await?;
        Ok(rows)*/

         // Get a client from the pool
        let client = self.pool.get().await?;
        
        // Execute the query and let the client be dropped automatically afterward
        let rows = client.query(query, params).await?;
        
        Ok(rows)

    } 
 

    pub async fn query_one(
        & self,
        query: &str,
        params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
    ) -> Result<tokio_postgres::Row, DatabaseError> {
            /*
          let client = &self.connect().await?;

        let rows =  client.query_one(query, params).await?;
        Ok(rows)

        */


         let client = self.pool.get().await ?;
        
        // Execute the query and let the client be dropped automatically afterward
        let row = client.query_one(query, params).await?;
        
        Ok(row)
    }
 
    //use for insert, update, etc
    pub async fn execute(
        &  self,
        query: &str,
        params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
    ) -> Result<u64, DatabaseError> {
        /*
          let client = &self.connect().await?;

        

        let rows = client.execute(query, params).await?;
        Ok(rows)*/


         // Get a client from the pool
        let client = self.pool.get().await?;
        
        // Execute the query and let the client be dropped automatically afterward
        let count = client.execute(query, params).await?;
        
        Ok(count)


    }   



    pub async fn check_connection(&self) -> Result<bool, DatabaseError> {
    // Get a client from the pool
    let client = self.pool.get().await?;
    
    // Execute a simple query to check the connection
    match client.execute("SELECT 1", &[]).await {
        Ok(_) => Ok(true),
        Err(e) => Err(DatabaseError::from(e))
    }
}

    pub async fn recreate_pool(&mut self) -> Result<(), DatabaseError> {
        // Parse the connection config from the URL
        let config: tokio_postgres::Config = self.connection_url.parse()
            .map_err(|_e| DatabaseError::ConnectionFailed)?;
            
        // Create a manager using the config
        let manager = deadpool_postgres::Manager::new(config, tokio_postgres::NoTls);
        
        // Create a new pool
        let new_pool = deadpool_postgres::Pool::builder(manager)
            .max_size(16)
            .build()
            .map_err(|e| DatabaseError::PoolCreationFailed(e.to_string()))?;
        
        // Replace the old pool
        self.pool = new_pool;
        
        Ok(())
    }

         
     pub async fn query_with_timeout(
        &self,
        query: &str,
        params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
    ) -> Result<Vec<tokio_postgres::Row>, DatabaseError> {
        match timeout(self.timeout_duration, self.query(query, params)).await {
            Ok(result) => result,
            Err(_) => Err(DatabaseError::ConnectionFailed) // Timeout occurred
        }
    }

    pub async fn query_with_retry(
        &self,
        query: &str,
        params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
    ) -> Result<Vec<tokio_postgres::Row>, DatabaseError> {
        let mut attempts = 0;
        
        while attempts < self.max_reconnect_attempts {
            match self.query(query, params).await {
                Ok(result) => return Ok(result),
                Err(e) => {
                    // If it's a connection error, retry
                    if let DatabaseError::PoolError(_) = e {
                        attempts += 1;
                        if attempts >= self.max_reconnect_attempts {
                            return Err(e);
                        }
                        // Wait before retrying
                        sleep(Duration::from_secs(2_u64.pow(attempts))).await;
                    } else {
                        // For other errors, return immediately
                        return Err(e);
                    }
                }
            }
        }
        
        Err(DatabaseError::ConnectionFailed)
    }


    /*async fn atomic_transaction(
        &  self,
        steps: Vec<PostgresInput<'_>>,
    ) -> Result<(), PostgresError> {
          let client = &mut self.connect().await?;

        // Start a transaction
        let transaction = client.transaction().await?;

        //for each step in steps
        for step in steps {
            //execute the step
            let result = transaction.execute(&step.query, step.params).await;
            //check if the result is ok
            if result.is_err() {
                //if not rollback
                transaction.rollback().await?;
                //return error
                return Err(PostgresError::from(result.err().unwrap()));
            }
        }

        //if all steps are ok commit
        transaction.commit().await?;
        //return ok
        Ok(())
    }*/
}

pub fn try_get_option<'a, T: tokio_postgres::types::FromSql<'a>>(
    row: &'a tokio_postgres::Row,
    column: &str,
) -> Option<T> {
    match row.try_get::<&str, T>(column) {
        Ok(value) => Some(value),
        Err(_) => None,
    }
}