pg_setup/
db.rs

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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
use crate::{db_cmd_strategy::CmdStrategy, error::Result, table_builder::TableBuilder};

/// Trait is used to implement actions of [`PostgresDB`].
#[async_trait]
pub(crate) trait DBStrategy {
    fn setup(&self, db_uri: &str) -> Result<()>;
    fn tear_down(&self, db_uri: &str) -> Result<()>;
    async fn execute(&self, db_uri: &str, sql: &str) -> Result<()>;
}

/// Builder to construct a [`PostgresDB`].
pub struct PostgresDBBuilder {
    db_uri: String,
    keep_db: bool,
    #[cfg(feature = "sqlx")]
    use_sqlx: bool,
    schema: String,
}

impl PostgresDBBuilder {
    pub fn new(db_uri: impl ToString) -> Self {
        Self {
            db_uri: db_uri.to_string(),
            keep_db: false,
            #[cfg(feature = "sqlx")]
            use_sqlx: false,
            schema: "public".to_string(),
        }
    }

    /// When set, does not drop the DB when [`PostgresDB`] goes out of scope.
    #[must_use]
    pub fn keep_db(mut self) -> Self {
        self.keep_db = true;
        self
    }

    /// The postgres schema to use. Defaults to `"public"`.
    #[must_use]
    pub fn schema(mut self, schema: impl ToString) -> Self {
        self.schema = schema.to_string();
        self
    }

    /// Use sqlx to create / drop the database? Requires that sqlx-cli is installed
    /// and in `$PATH` and that migrations are defined. Defaults to false.
    #[must_use]
    #[cfg(feature = "sqlx")]
    pub fn use_sqlx(mut self) -> Self {
        self.use_sqlx = true;
        self
    }

    pub async fn start(self) -> Result<PostgresDB> {
        #[cfg(feature = "sqlx")]
        let strategy: Box<dyn DBStrategy> = if self.use_sqlx {
            Box::new(crate::db_sqlx_strategy::SqlxStrategy)
        } else {
            Box::new(CmdStrategy)
        };
        #[cfg(not(feature = "sqlx"))]
        let strategy = Box::new(CmdStrategy);
        let db = PostgresDB {
            db_uri: self.db_uri,
            keep_db: self.keep_db,
            schema: self.schema,
            strategy,
        };
        db.setup().await?;
        Ok(db)
    }
}

/// See module comment.
pub struct PostgresDB {
    db_uri: String,
    keep_db: bool,
    schema: String,
    strategy: Box<dyn DBStrategy>,
}

impl PostgresDB {
    #[cfg(feature = "sqlx")]
    pub async fn con(&mut self) -> Result<sqlx::PgConnection> {
        crate::db_sqlx_strategy::connect(&self.db_uri).await
    }

    #[cfg(feature = "sqlx")]
    pub async fn pool(&self) -> Result<sqlx::Pool<sqlx::Postgres>> {
        crate::db_sqlx_strategy::pool(&self.db_uri).await
    }

    #[cfg(feature = "sqlx")]
    pub async fn delete_all_tables(&self) -> Result<()> {
        // use sqlx::postgres::PgPoolOptions;
        // use sqlx::{postgres::PgConnectOptions, prelude::*};

        use sqlx::Acquire;

        info!("WILL DELETE ALL TABLES");

        let pool = self.pool().await?;
        let mut t = pool.begin().await?;
        let con = t.acquire().await?;

        sqlx::query("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE")
            .execute(&mut *con)
            .await?;

        let tables: Vec<(String,)> = sqlx::query_as(
            "
    SELECT table_name
    FROM information_schema.tables
    WHERE table_schema LIKE $1 AND table_type = 'BASE TABLE';",
        )
        .bind(&self.schema)
        .fetch_all(&mut *con)
        .await?;

        for (table,) in tables {
            sqlx::query(&(format!("delete from {table};")))
                .execute(&mut *con)
                .await?;
        }
        t.commit().await?;

        Ok(())
    }

    pub async fn execute(&self, sql: &str) -> Result<()> {
        self.strategy.execute(&self.db_uri, sql).await
    }

    pub async fn create_table<F>(&self, table_name: impl AsRef<str>, tbl_callback: F) -> Result<()>
    where
        F: Fn(&mut TableBuilder),
    {
        let mut table = TableBuilder::new(format!("{}.{}", self.schema, table_name.as_ref()));
        tbl_callback(&mut table);
        let sql = table.to_sql()?;
        self.execute(sql.as_str()).await
    }

    pub async fn setup(&self) -> Result<()> {
        self.strategy.setup(&self.db_uri)
    }

    fn tear_down(&self) -> Result<()> {
        if self.keep_db {
            return Ok(());
        }

        self.strategy.tear_down(&self.db_uri)
    }
}

impl Drop for PostgresDB {
    fn drop(&mut self) {
        self.tear_down().expect("teardown");
    }
}