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

Code before:
        .map_err(|e| ConnectionError::BadConnection(e.to_string()))?;

    // The `conn` object performs actual IO with the database, and tokio-postgres suggests spawning
    // it off to run in the background. This will resolve only when the connection is closed, either
    // because of a fatal error or because its associated Client has dropped and all outstanding
    // work has completed.
    let conn_id = Uuid::new_v4();
    tokio::spawn(async move {
        if let Err(e) = conn.await {
            error!(%conn_id, "Database connection terminated: {e}");
        }
    });

    // Users interact with the database through the client object. We convert it into an
    // AsyncPgConnection so it can be compatible with diesel.
    match AsyncPgConnection::try_from(client).await {
        Ok(mut c) => {
            #[derive(QueryableByName)]
            struct Pid {
                #[diesel(sql_type = diesel::sql_types::Integer)]
                pg_backend_pid: i32,
            }

            let pid = diesel::sql_query("SELECT pg_backend_pid()")
                .get_result::<Pid>(&mut c)
                .await
                .map_err(ConnectionError::CouldntSetupConfiguration)?;

            info!(%conn_id, pid = ?pid.pg_backend_pid, "Database connection established");
            Ok(AsyncPgConnectionWithId {
                id: conn_id,
                conn: c,
            })
        }
        Err(e) => Err(e),
    }
}

/// Builds a TLS configuration from the provided DbArgs. If tls_verify_cert is false, disable server
/// certificate verification. If tls_ca_cert_path is provided, add the custom CA certificate to the
/// root certificates.
pub(crate) fn build_tls_config(
    tls_verify_cert: bool,
    tls_ca_cert_path: Option<PathBuf>,
) -> anyhow::Result<ClientConfig> {
    if !tls_verify_cert {
        return Ok(ClientConfig::builder()
            .dangerous()
            .with_custom_certificate_verifier(Arc::new(SkipServerCertCheck))
            .with_no_client_auth());
    }

Review comment (bug):
- Do you also want a "Database connection terminated" message when there isn't an error?
- You don't need the `match` here, you can keep using `?`
- I'm not fully following why you need the -- yes the `pid` could be re-used once a connection is dropped, but there should only be one PID per run, right? What about if you did the following:

```rust
let tls = MakeRustlsConnect::new(tls_config);
let (client, tokio_conn) = tokio_postgres::connect(database_url, tls)
    .await
    .map_err(|e| ConnectionError::BadConnection(e.to_string()))?;

let conn = AsyncPgConnection::try_from(client).await?;

#[derive(QueryableByName)]
struct Pid {
    #[diesel(sql_type = Integer)]
    pid: i32
}

let Pid { pid } = diesel::sql_query("SELECT pg_backend_pid() AS pid")
    .get_result(&mut c)
    .await
    .map_err(ConnectionError::CouldntSetupConfiguration)?;

tokio::spawn(async move {
    if let Err(e) = tokio_conn.await {
        error!(pid, "Database connection terminated with error: {e}");
    } else {
        debug!(pid, "Database connection terminated");
    }
});

Ok(AsyncConnectionWithPid { conn, pid })
```

Code after:
/// - If server doesn't support TLS: falls back to plaintext (when using sslmode=prefer default)
/// - If sslmode=disable in URL: no TLS attempted regardless of connector
pub(crate) async fn establish_tls_connection(
    database_url: &str,
    tls_config: ClientConfig,
) -> ConnectionResult<AsyncPgConnectionWithId> {
    let tls = MakeRustlsConnect::new(tls_config);
    let (client, tokio_conn) = tokio_postgres::connect(database_url, tls)
        .await
        .map_err(|e| ConnectionError::BadConnection(e.to_string()))?;

    // The `conn` object performs actual IO with the database, and tokio-postgres suggests spawning
    // it off to run in the background. This will resolve only when the connection is closed,
    // either because of a fatal error or because its associated Client has dropped and all
    // outstanding work has completed.
    let h_conn = tokio::spawn(tokio_conn);

    // Users interact with the database through the client object. We convert it into an
    // AsyncPgConnection so it can be compatible with diesel.
    let mut conn = AsyncPgConnection::try_from(client).await?;

    #[derive(QueryableByName)]
    struct Pid {
        #[diesel(sql_type = Integer)]
        pid: i32,
    }

    let Pid { pid } = diesel::sql_query("SELECT pg_backend_pid() AS pid")
        .get_result(&mut conn)
        .await
        .map_err(ConnectionError::CouldntSetupConfiguration)?;

    info!(pid, "Database connection established");

    tokio::spawn(async move {
        match h_conn.await {
            Ok(Err(e)) => {
                error!(pid, "Database connection terminated with error: {e}");
            }
            Ok(Ok(())) => {
                debug!(pid, "Database connection terminated without error");
            }
            Err(e) => {
                debug!(pid, "Database connection task failed: {e}");
            }
        }
    });

    Ok(AsyncPgConnectionWithId { conn, pid })
}