holdon 0.3.0

Wait for anything. Know why if it doesn't.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use std::sync::Arc;
use std::sync::OnceLock;

use rustls::ClientConfig;
use tokio_postgres::NoTls;
use tokio_postgres_rustls::MakeRustlsConnect;
use url::Url;

use super::hint::{Hintable, hints};
use super::{AttemptCtx, run_stage};
use crate::diagnostic::{Stage, StageKind};

pub(super) async fn probe(url: &Url, expect_table: Option<&str>, ctx: AttemptCtx) -> Vec<Stage> {
    let pw = url.password().unwrap_or("").to_owned();
    let want_tls = !sslmode_disabled(url);
    let driver_url = super::strip_query_keys(url, &["table"]);
    let driver_str = driver_url.as_str().to_owned();
    let table = expect_table.map(str::to_owned);
    vec![
        run_stage(
            StageKind::Postgres,
            ctx.attempt_timeout,
            hints::PG_NOT_READY,
            connect_and_query(driver_str.clone(), want_tls, table),
            &[driver_str.as_str(), pw.as_str()],
        )
        .await,
    ]
}

fn sslmode_disabled(url: &Url) -> bool {
    url.query_pairs()
        .any(|(k, v)| k.eq_ignore_ascii_case("sslmode") && v.eq_ignore_ascii_case("disable"))
}

#[derive(Debug, thiserror::Error)]
enum ProbeError {
    #[error(transparent)]
    Driver(#[from] tokio_postgres::Error),
    #[error("expected table `{0}` not found in information_schema.tables")]
    TableMissing(String),
}

impl Hintable for ProbeError {
    fn hint(&self) -> Option<&'static str> {
        match self {
            Self::Driver(e) => e.hint(),
            Self::TableMissing(_) => Some(hints::PG_TABLE_MISSING),
        }
    }
}

fn rustls_config() -> Arc<ClientConfig> {
    static CFG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
    CFG.get_or_init(|| {
        let _ = rustls::crypto::ring::default_provider().install_default();
        let mut roots = rustls::RootCertStore::empty();
        roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
        Arc::new(
            ClientConfig::builder()
                .with_root_certificates(roots)
                .with_no_client_auth(),
        )
    })
    .clone()
}

async fn connect_and_query(
    conn_str: String,
    want_tls: bool,
    expect_table: Option<String>,
) -> Result<(), ProbeError> {
    if want_tls {
        let tls = MakeRustlsConnect::new(rustls_config().as_ref().clone());
        let (client, connection) = tokio_postgres::connect(&conn_str, tls).await?;
        let driver = tokio::spawn(connection);
        let outcome = run_queries(&client, expect_table.as_deref()).await;
        drop(client);
        let _ = driver.await;
        outcome
    } else {
        let (client, connection) = tokio_postgres::connect(&conn_str, NoTls).await?;
        let driver = tokio::spawn(connection);
        let outcome = run_queries(&client, expect_table.as_deref()).await;
        drop(client);
        let _ = driver.await;
        outcome
    }
}

async fn run_queries(
    client: &tokio_postgres::Client,
    expect_table: Option<&str>,
) -> Result<(), ProbeError> {
    client.simple_query("SELECT 1").await?;
    if let Some(name) = expect_table {
        let row = client
            .query_opt(
                "SELECT 1 FROM information_schema.tables \
                 WHERE table_name = $1 AND table_schema = ANY (current_schemas(false)) \
                 LIMIT 1",
                &[&name],
            )
            .await?;
        if row.is_none() {
            return Err(ProbeError::TableMissing(name.to_owned()));
        }
    }
    Ok(())
}