Skip to main content

meta_srv/utils/
postgres.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#[cfg(unix)]
16use std::str::FromStr;
17
18use common_error::ext::BoxedError;
19use common_meta::election::ElectionRef;
20use common_meta::election::rds::postgres::{ElectionPgClient, PgElection};
21use common_meta::kv_backend::KvBackendRef;
22use common_meta::kv_backend::rds::PgStore;
23use common_meta::kv_backend::rds::postgres::{
24    TlsMode as PgTlsMode, TlsOption as PgTlsOption, create_postgres_tls_connector,
25};
26use common_telemetry::warn;
27use deadpool_postgres::{Config, Runtime};
28use servers::tls::TlsOption;
29use snafu::{OptionExt, ResultExt};
30use tokio_postgres::NoTls;
31#[cfg(unix)]
32use tokio_postgres::config::Host;
33
34use crate::error::{self, Result};
35
36/// Converts [`TlsOption`] to [`PgTlsOption`] to avoid circular dependencies
37fn convert_tls_option(tls_option: &TlsOption) -> PgTlsOption {
38    let mode = match tls_option.mode {
39        servers::tls::TlsMode::Disable => PgTlsMode::Disable,
40        servers::tls::TlsMode::Prefer => PgTlsMode::Prefer,
41        servers::tls::TlsMode::Require => PgTlsMode::Require,
42        servers::tls::TlsMode::VerifyCa => PgTlsMode::VerifyCa,
43        servers::tls::TlsMode::VerifyFull => PgTlsMode::VerifyFull,
44    };
45
46    PgTlsOption {
47        mode,
48        cert_path: tls_option.cert_path.clone(),
49        key_path: tls_option.key_path.clone(),
50        ca_cert_path: tls_option.ca_cert_path.clone(),
51        watch: tls_option.watch,
52    }
53}
54
55/// Creates a pool for the Postgres backend with config and optional TLS.
56///
57/// It only use first store addr to create a pool, and use the given config to create a pool.
58pub async fn create_postgres_pool(
59    store_addrs: &[String],
60    cfg: Option<Config>,
61    tls_config: Option<TlsOption>,
62) -> Result<deadpool_postgres::Pool> {
63    let mut cfg = cfg.unwrap_or_default();
64    let postgres_url = store_addrs.first().context(error::InvalidArgumentsSnafu {
65        err_msg: "empty store addrs",
66    })?;
67    cfg.url = Some(postgres_url.clone());
68
69    let is_unix_socket = is_unix_socket_url(postgres_url);
70    if is_unix_socket
71        && matches!(tls_config.as_ref(), Some(t) if t.mode != servers::tls::TlsMode::Disable)
72    {
73        warn!(
74            "TLS is not supported for Unix domain socket PostgreSQL connections, falling back to NoTls"
75        );
76    }
77
78    let pool = match tls_config {
79        Some(tls_config)
80            if tls_config.mode != servers::tls::TlsMode::Disable && !is_unix_socket =>
81        {
82            let pg_tls_config = convert_tls_option(&tls_config);
83            let tls_connector =
84                create_postgres_tls_connector(&pg_tls_config).map_err(|e| error::Error::Other {
85                    source: BoxedError::new(e),
86                    location: snafu::Location::new(file!(), line!(), 0),
87                })?;
88            cfg.create_pool(Some(Runtime::Tokio1), tls_connector)
89                .context(error::CreatePostgresPoolSnafu)?
90        }
91        _ => cfg
92            .create_pool(Some(Runtime::Tokio1), NoTls)
93            .context(error::CreatePostgresPoolSnafu)?,
94    };
95
96    Ok(pool)
97}
98
99#[cfg(unix)]
100fn is_unix_socket_url(url: &str) -> bool {
101    let Ok(cfg) = tokio_postgres::Config::from_str(url) else {
102        return false;
103    };
104    // tokio-postgres dials `hostaddr` over TCP even when `host` is a socket path,
105    // so treat the config as a socket only when every host is a Unix socket and
106    // no `hostaddr` is set.
107    cfg.get_hostaddrs().is_empty()
108        && !cfg.get_hosts().is_empty()
109        && cfg
110            .get_hosts()
111            .iter()
112            .all(|host| matches!(host, Host::Unix(_)))
113}
114
115#[cfg(not(unix))]
116fn is_unix_socket_url(_: &str) -> bool {
117    false
118}
119
120/// Builds a Postgres-backed metadata [`KvBackendRef`].
121///
122/// * `store_addrs` - Postgres connection URLs; only the first address is used.
123/// * `cfg` - optional deadpool config to customize pool/session behavior.
124/// * `tls_config` - optional TLS settings for the Postgres connection.
125/// * `schema_name` - optional schema containing the metadata table.
126/// * `table_name` - metadata KV table name.
127/// * `max_txn_ops` - maximum operations allowed in one metadata transaction.
128/// * `auto_create_schema` - whether to create `schema_name` when it is missing.
129#[allow(clippy::too_many_arguments)]
130pub async fn build_postgres_kv_backend(
131    store_addrs: &[String],
132    cfg: Option<Config>,
133    tls_config: Option<TlsOption>,
134    schema_name: Option<&str>,
135    table_name: &str,
136    max_txn_ops: usize,
137    auto_create_schema: bool,
138) -> Result<KvBackendRef> {
139    let pool = create_postgres_pool(store_addrs, cfg, tls_config).await?;
140    PgStore::with_pg_pool(
141        pool,
142        schema_name,
143        table_name,
144        max_txn_ops,
145        auto_create_schema,
146    )
147    .await
148    .context(error::KvBackendSnafu)
149}
150
151/// Builds a Postgres-backed election implementation.
152///
153/// * `store_addrs` - Postgres connection URLs; only the first address is used.
154/// * `cfg` - optional deadpool config to customize pool/session behavior.
155/// * `tls_config` - optional TLS settings for the Postgres connection.
156/// * `leader_value` - advertised address of this election candidate.
157/// * `store_key_prefix` - prefix for election and candidate keys.
158/// * `candidate_lease_ttl` - TTL for registered candidate metadata.
159/// * `meta_lease_ttl` - TTL for the elected leader metadata.
160/// * `schema_name` - optional schema containing the metadata table.
161/// * `table_name` - metadata KV table name used for election records.
162/// * `lock_id` - Postgres advisory lock id used by the election.
163#[allow(clippy::too_many_arguments)]
164pub async fn build_postgres_election(
165    store_addrs: &[String],
166    cfg: Option<Config>,
167    tls_config: Option<TlsOption>,
168    leader_value: String,
169    store_key_prefix: String,
170    candidate_lease_ttl: std::time::Duration,
171    meta_lease_ttl: std::time::Duration,
172    schema_name: Option<&str>,
173    table_name: &str,
174    lock_id: u64,
175) -> Result<ElectionRef> {
176    let pool = create_postgres_pool(store_addrs, cfg, tls_config).await?;
177    let election_client =
178        ElectionPgClient::new(pool, meta_lease_ttl, meta_lease_ttl, meta_lease_ttl)
179            .context(error::KvBackendSnafu)?;
180    PgElection::with_pg_client(
181        leader_value,
182        election_client,
183        store_key_prefix,
184        candidate_lease_ttl,
185        meta_lease_ttl,
186        schema_name,
187        table_name,
188        lock_id,
189    )
190    .await
191    .context(error::KvBackendSnafu)
192}
193
194#[cfg(test)]
195mod tests {
196    use super::is_unix_socket_url;
197
198    #[test]
199    fn detects_postgres_unix_socket_url() {
200        #[cfg(unix)]
201        {
202            // libpq keyword-value form (issue #7734)
203            assert!(is_unix_socket_url(
204                "host=/var/run/postgresql dbname=greptime user=greptime password=secret"
205            ));
206            // standard postgres URL with percent-encoded unix socket directory
207            assert!(is_unix_socket_url(
208                "postgresql://user:pw@%2Fvar%2Frun%2Fpostgresql/mydb"
209            ));
210            assert!(is_unix_socket_url(
211                "postgresql://user@%2Fvar%2Frun%2Fpostgresql/db"
212            ));
213            // postgres URL with socket dir in query param
214            assert!(is_unix_socket_url(
215                "postgresql:///mydb?host=%2Fvar%2Frun%2Fpostgresql"
216            ));
217            assert!(is_unix_socket_url(
218                "postgresql://user:secret@/mydb?host=%2Fvar%2Frun%2Fpostgresql"
219            ));
220
221            // TCP URLs should not be classified as unix socket
222            assert!(!is_unix_socket_url(
223                "postgresql://user:pw@localhost:5432/mydb"
224            ));
225            assert!(!is_unix_socket_url("postgresql://user@localhost/db"));
226            assert!(!is_unix_socket_url(
227                "host=127.0.0.1 port=5432 dbname=greptime user=greptime password=secret"
228            ));
229            // mixed socket and TCP hosts must not be treated as a socket-only config:
230            // a Require/VerifyFull TLS config would otherwise end up in plaintext over TCP
231            assert!(!is_unix_socket_url(
232                "host=/var/run/postgresql,db.example.com"
233            ));
234            // tokio-postgres dials `hostaddr` over TCP even when `host` is a socket path
235            assert!(!is_unix_socket_url("host=/tmp hostaddr=10.0.0.5"));
236        }
237
238        #[cfg(not(unix))]
239        {
240            assert!(!is_unix_socket_url(
241                "host=/var/run/postgresql dbname=greptime user=greptime password=secret"
242            ));
243            assert!(!is_unix_socket_url(
244                "postgresql://user:pw@localhost:5432/mydb"
245            ));
246        }
247    }
248}