wb_cache/test/simulation/db/driver/
pg.rs

1//! PostgreSQL database driver.
2use std::time::Duration;
3
4use async_trait::async_trait;
5use fieldx::fxstruct;
6use sea_orm::ConnectOptions;
7use sea_orm::ConnectionTrait;
8use sea_orm::DatabaseConnection;
9
10use crate::test::simulation::types::Result;
11
12use super::DatabaseDriver;
13
14/// PostgreSQL database driver.
15#[derive(Debug)]
16#[fxstruct(sync, rc, no_new, builder)]
17pub struct Pg {
18    host:       String,
19    port:       u16,
20    user:       String,
21    password:   String,
22    database:   String,
23    #[fieldx(inner_mut, get(off), set, builder(off))]
24    connection: DatabaseConnection,
25}
26
27impl Pg {
28    pub async fn connect(&self) -> Result<()> {
29        let schema = format!(
30            "postgres://{}:{}@{}:{}/{}",
31            self.user, self.password, self.host, self.port, self.database
32        );
33        let mut opts = ConnectOptions::new(&schema);
34        opts.max_connections(20)
35            .acquire_timeout(Duration::from_secs(10))
36            .idle_timeout(Duration::from_secs(20))
37            .max_lifetime(Duration::from_secs(60))
38            .test_before_acquire(true);
39
40        self.set_connection(
41            sea_orm::Database::connect(opts)
42                .await
43                .inspect_err(|e| eprintln!("Error connecting to database {schema}: {e}"))?,
44        );
45
46        Ok(())
47    }
48}
49
50#[async_trait]
51impl DatabaseDriver for Pg {
52    fn connection(&self) -> DatabaseConnection {
53        self.connection.read().clone()
54    }
55
56    /// Turn off synchronous commit for performance.
57    async fn configure(&self) -> Result<()> {
58        self.connection()
59            .execute_unprepared("SET synchronous_commit = off;")
60            .await?;
61
62        Ok(())
63    }
64
65    async fn checkpoint(&self) -> Result<()> {
66        Ok(())
67    }
68}