elefant-tools 0.0.2

A library for doing things like pg_dump and pg_restore, with extra features, and probably more bugs.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

use crate::{AsyncCleanup, PostgresClientWrapper};
use std::sync::Arc;
use tokio::sync::Mutex;

/// A very simple connection pool for Postgres connections.
#[derive(Clone, Default)]
pub struct ConnectionPool {
    connection_pool: Arc<Mutex<Vec<PostgresClientWrapper>>>,
}

impl ConnectionPool {
    pub fn new() -> Self {
        Self {
            connection_pool: Arc::new(Mutex::new(Vec::new())),
        }
    }

    /// Get a connection from the pool. If the pool is empty, this will return `None`.
    pub async fn get_connection(&self) -> Option<PostgresClientWrapper> {
        let mut pool = self.connection_pool.lock().await;
        pool.pop()
    }

    /// Release a connection back to the pool.
    pub async fn release_connection(&self, connection: PostgresClientWrapper) {
        let mut pool = self.connection_pool.lock().await;
        pool.push(connection);
    }
}

/// A wrapper around a connection that will release the connection back to the pool when `cleanup`
/// is called.
pub struct ReleaseConnection {
    pool: ConnectionPool,
    connection: PostgresClientWrapper,
}

impl ReleaseConnection {
    /// Create a new `ReleaseConnection` instance.
    pub fn new(pool: ConnectionPool, connection: PostgresClientWrapper) -> Self {
        Self { pool, connection }
    }
}

impl AsyncCleanup for ReleaseConnection {
    async fn cleanup(self) -> crate::Result<()> {
        self.pool.release_connection(self.connection).await;
        Ok(())
    }
}