Skip to main content

refinery_core/drivers/
tokio_postgres.rs

1use crate::traits::r#async::{AsyncMigrate, AsyncQuery, AsyncTransaction};
2use crate::Migration;
3use time::format_description::well_known::Rfc3339;
4use time::OffsetDateTime;
5use tokio_postgres::error::Error as PgError;
6use tokio_postgres::{Client, Transaction as PgTransaction};
7
8async fn query_applied_migrations(
9    transaction: &PgTransaction<'_>,
10    query: &str,
11) -> Result<Vec<Migration>, PgError> {
12    let rows = transaction.query(query, &[]).await?;
13    let mut applied = Vec::new();
14    for row in rows.into_iter() {
15        let version = row.get(0);
16        let applied_on: String = row.get(2);
17        // Safe to call unwrap, as we stored it in RFC3339 format on the database
18        let applied_on = OffsetDateTime::parse(&applied_on, &Rfc3339).unwrap();
19        let checksum: String = row.get(3);
20
21        applied.push(Migration::applied(
22            version,
23            row.get(1),
24            applied_on,
25            checksum
26                .parse::<u64>()
27                .expect("checksum must be a valid u64"),
28        ));
29    }
30    Ok(applied)
31}
32
33impl AsyncTransaction for Client {
34    type Error = PgError;
35
36    async fn execute<'a, T: Iterator<Item = &'a str> + Send + 'a>(
37        &mut self,
38        queries: T,
39    ) -> Result<usize, Self::Error> {
40        let transaction = self.transaction().await?;
41        let mut count = 0;
42        for query in queries {
43            transaction.batch_execute(query).await?;
44            count += 1;
45        }
46        transaction.commit().await?;
47        Ok(count as usize)
48    }
49}
50
51impl AsyncQuery<Vec<Migration>> for Client {
52    async fn query(
53        &mut self,
54        query: &str,
55    ) -> Result<Vec<Migration>, <Self as AsyncTransaction>::Error> {
56        let transaction = self.transaction().await?;
57        let applied = query_applied_migrations(&transaction, query).await?;
58        transaction.commit().await?;
59        Ok(applied)
60    }
61}
62
63impl AsyncMigrate for Client {}