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

Code before:
use prometheus::core::GenericCounter;
use tracing::info;
use url::Url;

use crate::tls::AsyncPgConnectionWithId;
use crate::tls::build_tls_config;
use crate::tls::establish_tls_connection;

mod model;
pub mod query;
pub mod schema;
pub mod store;
pub mod temp;
mod tls;

pub use sui_field_count::FieldCount;
pub use sui_sql_macro::sql;

pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");

#[derive(Debug, Clone)]
pub struct DbConfig<'a> {
    database_url: Url,
    db_args: DbArgs,
    read_only: bool,
    registry: Option<&'a Registry>,
}

#[derive(clap::Args, Debug, Clone)]
pub struct DbArgs {
    /// Number of connections to keep in the pool.
    #[arg(long, default_value_t = Self::default().db_connection_pool_size)]
    pub db_connection_pool_size: u32,

    /// Time spent waiting for a connection from the pool to become available, in milliseconds.
    #[arg(long, default_value_t = Self::default().db_connection_timeout_ms)]
    pub db_connection_timeout_ms: u64,

    #[arg(long)]
    /// Time spent waiting for statements to complete, in milliseconds.
    pub db_statement_timeout_ms: Option<u64>,

    #[arg(long)]
    /// Enable server certificate verification. By default, this is set to false to match the
    /// default behavior of libpq.
    pub tls_verify_cert: bool,

    #[arg(long)]
    /// Path to a custom CA certificate to use for server certificate verification.
    pub tls_ca_cert_path: Option<PathBuf>,
}

Review comment (suggestion):
It is a bit confusing to introduce the prometheus registry here -- everywhere else `Config` types are plain old data.

It looks like the reason to do this is to offer an API, `Db::register_stats(&Registry)` to optionally fill out the `PoolStats`. That way you can avoid updating all the `Db` constructor call sites as well.

Code after:

use crate::tls::AsyncPgConnectionWithId;
use crate::tls::build_tls_config;
use crate::tls::establish_tls_connection;

mod metrics;
mod model;
pub mod query;
pub mod schema;
pub mod store;
pub mod temp;
mod tls;

use crate::metrics::PoolMetrics;
pub use sui_field_count::FieldCount;
pub use sui_sql_macro::sql;

pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");

#[derive(clap::Args, Debug, Clone)]
pub struct DbArgs {
    /// Number of connections to keep in the pool.
    #[arg(long, default_value_t = Self::default().db_connection_pool_size)]
    pub db_connection_pool_size: u32,

    /// Time spent waiting for a connection from the pool to become available, in milliseconds.
    #[arg(long, default_value_t = Self::default().db_connection_timeout_ms)]
    pub db_connection_timeout_ms: u64,

    #[arg(long)]
    /// Time spent waiting for statements to complete, in milliseconds.
    pub db_statement_timeout_ms: Option<u64>,

    #[arg(long)]
    /// Enable server certificate verification. By default, this is set to false to match the
    /// default behavior of libpq.
    pub tls_verify_cert: bool,

    #[arg(long)]
    /// Path to a custom CA certificate to use for server certificate verification.
    pub tls_ca_cert_path: Option<PathBuf>,
}

#[derive(Clone)]
pub struct Db {
    pool: Pool<AsyncPgConnectionWithId>,
    pool_metrics: Option<Arc<PoolMetrics>>,
}

// Ensures that unacquired_canceled is incremented if neither acquired nor unacquired_error are so
// that one of acquired, unacquired_error, unacquired_canceled is incremented for every requested