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
15use common_error::ext::BoxedError;
16use common_meta::election::ElectionRef;
17use common_meta::election::rds::postgres::{ElectionPgClient, PgElection};
18use common_meta::kv_backend::KvBackendRef;
19use common_meta::kv_backend::rds::PgStore;
20use common_meta::kv_backend::rds::postgres::{
21    TlsMode as PgTlsMode, TlsOption as PgTlsOption, create_postgres_tls_connector,
22};
23use deadpool_postgres::{Config, Runtime};
24use servers::tls::TlsOption;
25use snafu::{OptionExt, ResultExt};
26use tokio_postgres::NoTls;
27
28use crate::error::{self, Result};
29
30/// Converts [`TlsOption`] to [`PgTlsOption`] to avoid circular dependencies
31fn convert_tls_option(tls_option: &TlsOption) -> PgTlsOption {
32    let mode = match tls_option.mode {
33        servers::tls::TlsMode::Disable => PgTlsMode::Disable,
34        servers::tls::TlsMode::Prefer => PgTlsMode::Prefer,
35        servers::tls::TlsMode::Require => PgTlsMode::Require,
36        servers::tls::TlsMode::VerifyCa => PgTlsMode::VerifyCa,
37        servers::tls::TlsMode::VerifyFull => PgTlsMode::VerifyFull,
38    };
39
40    PgTlsOption {
41        mode,
42        cert_path: tls_option.cert_path.clone(),
43        key_path: tls_option.key_path.clone(),
44        ca_cert_path: tls_option.ca_cert_path.clone(),
45        watch: tls_option.watch,
46    }
47}
48
49/// Creates a pool for the Postgres backend with config and optional TLS.
50///
51/// It only use first store addr to create a pool, and use the given config to create a pool.
52pub async fn create_postgres_pool(
53    store_addrs: &[String],
54    cfg: Option<Config>,
55    tls_config: Option<TlsOption>,
56) -> Result<deadpool_postgres::Pool> {
57    let mut cfg = cfg.unwrap_or_default();
58    let postgres_url = store_addrs.first().context(error::InvalidArgumentsSnafu {
59        err_msg: "empty store addrs",
60    })?;
61    cfg.url = Some(postgres_url.clone());
62
63    let pool = if let Some(tls_config) = tls_config {
64        let pg_tls_config = convert_tls_option(&tls_config);
65        let tls_connector =
66            create_postgres_tls_connector(&pg_tls_config).map_err(|e| error::Error::Other {
67                source: BoxedError::new(e),
68                location: snafu::Location::new(file!(), line!(), 0),
69            })?;
70        cfg.create_pool(Some(Runtime::Tokio1), tls_connector)
71            .context(error::CreatePostgresPoolSnafu)?
72    } else {
73        cfg.create_pool(Some(Runtime::Tokio1), NoTls)
74            .context(error::CreatePostgresPoolSnafu)?
75    };
76
77    Ok(pool)
78}
79
80/// Builds a Postgres-backed metadata [`KvBackendRef`].
81///
82/// * `store_addrs` - Postgres connection URLs; only the first address is used.
83/// * `cfg` - optional deadpool config to customize pool/session behavior.
84/// * `tls_config` - optional TLS settings for the Postgres connection.
85/// * `schema_name` - optional schema containing the metadata table.
86/// * `table_name` - metadata KV table name.
87/// * `max_txn_ops` - maximum operations allowed in one metadata transaction.
88/// * `auto_create_schema` - whether to create `schema_name` when it is missing.
89#[allow(clippy::too_many_arguments)]
90pub async fn build_postgres_kv_backend(
91    store_addrs: &[String],
92    cfg: Option<Config>,
93    tls_config: Option<TlsOption>,
94    schema_name: Option<&str>,
95    table_name: &str,
96    max_txn_ops: usize,
97    auto_create_schema: bool,
98) -> Result<KvBackendRef> {
99    let pool = create_postgres_pool(store_addrs, cfg, tls_config).await?;
100    PgStore::with_pg_pool(
101        pool,
102        schema_name,
103        table_name,
104        max_txn_ops,
105        auto_create_schema,
106    )
107    .await
108    .context(error::KvBackendSnafu)
109}
110
111/// Builds a Postgres-backed election implementation.
112///
113/// * `store_addrs` - Postgres connection URLs; only the first address is used.
114/// * `cfg` - optional deadpool config to customize pool/session behavior.
115/// * `tls_config` - optional TLS settings for the Postgres connection.
116/// * `leader_value` - advertised address of this election candidate.
117/// * `store_key_prefix` - prefix for election and candidate keys.
118/// * `candidate_lease_ttl` - TTL for registered candidate metadata.
119/// * `meta_lease_ttl` - TTL for the elected leader metadata.
120/// * `schema_name` - optional schema containing the metadata table.
121/// * `table_name` - metadata KV table name used for election records.
122/// * `lock_id` - Postgres advisory lock id used by the election.
123#[allow(clippy::too_many_arguments)]
124pub async fn build_postgres_election(
125    store_addrs: &[String],
126    cfg: Option<Config>,
127    tls_config: Option<TlsOption>,
128    leader_value: String,
129    store_key_prefix: String,
130    candidate_lease_ttl: std::time::Duration,
131    meta_lease_ttl: std::time::Duration,
132    schema_name: Option<&str>,
133    table_name: &str,
134    lock_id: u64,
135) -> Result<ElectionRef> {
136    let pool = create_postgres_pool(store_addrs, cfg, tls_config).await?;
137    let election_client =
138        ElectionPgClient::new(pool, meta_lease_ttl, meta_lease_ttl, meta_lease_ttl)
139            .context(error::KvBackendSnafu)?;
140    PgElection::with_pg_client(
141        leader_value,
142        election_client,
143        store_key_prefix,
144        candidate_lease_ttl,
145        meta_lease_ttl,
146        schema_name,
147        table_name,
148        lock_id,
149    )
150    .await
151    .context(error::KvBackendSnafu)
152}