pg_setup/
db.rs

1use crate::{db_cmd_strategy::CmdStrategy, error::Result, table_builder::TableBuilder};
2
3/// Trait is used to implement actions of [`PostgresDB`].
4#[async_trait]
5pub(crate) trait DBStrategy {
6    fn setup(&self, db_uri: &str) -> Result<()>;
7    fn tear_down(&self, db_uri: &str) -> Result<()>;
8    async fn execute(&self, db_uri: &str, sql: &str) -> Result<()>;
9}
10
11/// Builder to construct a [`PostgresDB`].
12pub struct PostgresDBBuilder {
13    db_uri: String,
14    keep_db: bool,
15    #[cfg(feature = "sqlx")]
16    use_sqlx: bool,
17    schema: String,
18}
19
20impl PostgresDBBuilder {
21    pub fn new(db_uri: impl ToString) -> Self {
22        Self {
23            db_uri: db_uri.to_string(),
24            keep_db: false,
25            #[cfg(feature = "sqlx")]
26            use_sqlx: false,
27            schema: "public".to_string(),
28        }
29    }
30
31    /// When set, does not drop the DB when [`PostgresDB`] goes out of scope.
32    #[must_use]
33    pub fn keep_db(mut self) -> Self {
34        self.keep_db = true;
35        self
36    }
37
38    /// The postgres schema to use. Defaults to `"public"`.
39    #[must_use]
40    pub fn schema(mut self, schema: impl ToString) -> Self {
41        self.schema = schema.to_string();
42        self
43    }
44
45    /// Use sqlx to create / drop the database? Requires that sqlx-cli is installed
46    /// and in `$PATH` and that migrations are defined. Defaults to false.
47    #[must_use]
48    #[cfg(feature = "sqlx")]
49    pub fn use_sqlx(mut self) -> Self {
50        self.use_sqlx = true;
51        self
52    }
53
54    pub async fn start(self) -> Result<PostgresDB> {
55        #[cfg(feature = "sqlx")]
56        let strategy: Box<dyn DBStrategy> = if self.use_sqlx {
57            Box::new(crate::db_sqlx_strategy::SqlxStrategy)
58        } else {
59            Box::new(CmdStrategy)
60        };
61        #[cfg(not(feature = "sqlx"))]
62        let strategy = Box::new(CmdStrategy);
63        let db = PostgresDB {
64            db_uri: self.db_uri,
65            keep_db: self.keep_db,
66            schema: self.schema,
67            strategy,
68        };
69        db.setup().await?;
70        Ok(db)
71    }
72}
73
74/// See module comment.
75pub struct PostgresDB {
76    db_uri: String,
77    keep_db: bool,
78    schema: String,
79    strategy: Box<dyn DBStrategy>,
80}
81
82impl PostgresDB {
83    #[cfg(feature = "sqlx")]
84    pub async fn con(&mut self) -> Result<sqlx::PgConnection> {
85        crate::db_sqlx_strategy::connect(&self.db_uri).await
86    }
87
88    #[cfg(feature = "sqlx")]
89    pub async fn pool(&self) -> Result<sqlx::Pool<sqlx::Postgres>> {
90        crate::db_sqlx_strategy::pool(&self.db_uri).await
91    }
92
93    #[cfg(feature = "sqlx")]
94    pub async fn delete_all_tables(&self) -> Result<()> {
95        // use sqlx::postgres::PgPoolOptions;
96        // use sqlx::{postgres::PgConnectOptions, prelude::*};
97
98        use sqlx::Acquire;
99
100        info!("WILL DELETE ALL TABLES");
101
102        let pool = self.pool().await?;
103        let mut t = pool.begin().await?;
104        let con = t.acquire().await?;
105
106        sqlx::query("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE")
107            .execute(&mut *con)
108            .await?;
109
110        let tables: Vec<(String,)> = sqlx::query_as(
111            "
112    SELECT table_name
113    FROM information_schema.tables
114    WHERE table_schema LIKE $1 AND table_type = 'BASE TABLE';",
115        )
116        .bind(&self.schema)
117        .fetch_all(&mut *con)
118        .await?;
119
120        for (table,) in tables {
121            sqlx::query(&(format!("delete from {table};")))
122                .execute(&mut *con)
123                .await?;
124        }
125        t.commit().await?;
126
127        Ok(())
128    }
129
130    pub async fn execute(&self, sql: &str) -> Result<()> {
131        self.strategy.execute(&self.db_uri, sql).await
132    }
133
134    pub async fn create_table<F>(&self, table_name: impl AsRef<str>, tbl_callback: F) -> Result<()>
135    where
136        F: Fn(&mut TableBuilder),
137    {
138        let mut table = TableBuilder::new(format!("{}.{}", self.schema, table_name.as_ref()));
139        tbl_callback(&mut table);
140        let sql = table.to_sql()?;
141        self.execute(sql.as_str()).await
142    }
143
144    pub async fn setup(&self) -> Result<()> {
145        self.strategy.setup(&self.db_uri)
146    }
147
148    fn tear_down(&self) -> Result<()> {
149        if self.keep_db {
150            return Ok(());
151        }
152
153        self.strategy.tear_down(&self.db_uri)
154    }
155}
156
157impl Drop for PostgresDB {
158    fn drop(&mut self) {
159        self.tear_down().expect("teardown");
160    }
161}
pg_setup/db.rs

pg_setup/
db.rs
