Repository: MystenLabs/sui
File: crates/sui-pg-db/src/lib.rs (Rust)

Code before:

impl DbArgs {
    pub fn connection_timeout(&self) -> Duration {
        Duration::from_millis(self.db_connection_timeout_ms)
    }

    pub fn statement_timeout(&self) -> Option<Duration> {
        self.db_statement_timeout_ms.map(Duration::from_millis)
    }
}

impl Db {
    /// Construct a new DB connection pool talking to the database at `database_url`.
    /// Instances of [Db] can be cloned to share access to the same pool
    pub async fn new(
        DbConfig {
            database_url,
            db_args,
            read_only,
            registry,
        }: DbConfig<'_>,
    ) -> anyhow::Result<Self> {
        Ok(Db {
            pool: pool(database_url, db_args, read_only).await?,
            pool_stats: registry.map(|registry| {
                let register = |name: &str, help: &str| -> GenericCounter<AtomicU64> {
                    let counter = IntCounter::new(name, help).unwrap();
                    registry.register(Box::new(counter.clone())).unwrap();
                    counter
                };
                Arc::new(PoolStats {
                    requested: register(
                        "connections_requested",
                        "Total requested connections from the pool",
                    ),
                    acquired: register(
                        "connections_acquired",
                        "Total requested connections from the pool that were acquired",
                    ),
                    unacquired_error: register(
                        "connections_unacquired_error",
                        "Total requested connections from the pool that were not acquired due to an error",
                    ),
                    unacquired_canceled: register(
                        "connections_unacquired_canceled",
                        "Total requested connections from the pool that were not acquired due to task cancellation",
                    ),
                })
            })
        })
    }

Review comment (refactor):
Let's align this with how other metrics sets are defined, for consistency/discoverability purposes:

- Can we use the macros supplied with `prometheus` instead of defining a new closure? (Check out how the other metrics structs are defined).
- Maybe call this `PoolMetrics` for consistency as well.
- ...and define a `PoolMetrics::new(&Registry): Arc<PoolMetrics>` 
- ...and perhaps move this into a `metrics` module for this crate.

Code after:
    pub async fn connect(&self) -> anyhow::Result<Connection<'_>> {
        if let Some(pool_metrics) = &self.pool_metrics {
            let guard = CancelGuard::request(pool_metrics);
            match self.pool.get().await {
                Ok(c) => {
                    guard.acquired();
                    Ok(Connection(c))
                }
                Err(e) => {
                    guard.unacquired_error();
                    Err(e.into())
                }
            }
        } else {
            Ok(Connection(self.pool.get().await?))
        }
    }

    /// Statistics about the connection pool
    pub fn state(&self) -> bb8::State {
        self.pool.state()
    }

    pub fn pool_metrics(&self) -> Option<&PoolMetrics> {
        self.pool_metrics.as_ref().map(|p| p.as_ref())
    }

    async fn clear_database(&self) -> anyhow::Result<()> {
        info!("Clearing the database...");
        let mut conn = self.connect().await?;
        let drop_all_tables = "
        DO $$ DECLARE
            r RECORD;
        BEGIN
        FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public')
            LOOP
                EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';
            END LOOP;
        END $$;";
        diesel::sql_query(drop_all_tables)
            .execute(&mut conn)
            .await?;
        info!("Dropped all tables.");

        let drop_all_procedures = "
        DO $$ DECLARE
            r RECORD;
        BEGIN
            FOR r IN (SELECT proname, oidvectortypes(proargtypes) as argtypes
                      FROM pg_proc INNER JOIN pg_namespace ns ON (pg_proc.pronamespace = ns.oid)
                      WHERE ns.nspname = 'public' AND prokind = 'p')