mz_postgres_util/
replication.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::str::FromStr;
11
12use mz_sql_parser::ast::{Ident, display::AstDisplay};
13use tokio_postgres::{
14    Client,
15    types::{Oid, PgLsn},
16};
17
18use mz_ssh_util::tunnel_manager::SshTunnelManager;
19
20use crate::{Config, PostgresError, simple_query_opt};
21
22#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
23pub enum WalLevel {
24    Minimal,
25    Replica,
26    Logical,
27}
28
29impl std::str::FromStr for WalLevel {
30    type Err = anyhow::Error;
31    fn from_str(s: &str) -> Result<Self, Self::Err> {
32        match s {
33            "minimal" => Ok(Self::Minimal),
34            "replica" => Ok(Self::Replica),
35            "logical" => Ok(Self::Logical),
36            o => Err(anyhow::anyhow!("unknown wal_level {}", o)),
37        }
38    }
39}
40
41impl std::fmt::Display for WalLevel {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        let s = match self {
44            WalLevel::Minimal => "minimal",
45            WalLevel::Replica => "replica",
46            WalLevel::Logical => "logical",
47        };
48
49        f.write_str(s)
50    }
51}
52
53#[mz_ore::test]
54fn test_wal_level_max() {
55    // Ensure `WalLevel::Logical` is the max among all levels.
56    for o in [WalLevel::Minimal, WalLevel::Replica, WalLevel::Logical] {
57        assert_eq!(WalLevel::Logical, WalLevel::Logical.max(o))
58    }
59}
60
61pub async fn get_wal_level(client: &Client) -> Result<WalLevel, PostgresError> {
62    let wal_level = client.query_one("SHOW wal_level", &[]).await?;
63    let wal_level: String = wal_level.get("wal_level");
64    Ok(WalLevel::from_str(&wal_level)?)
65}
66
67pub async fn get_max_wal_senders(client: &Client) -> Result<i64, PostgresError> {
68    let max_wal_senders = client
69        .query_one(
70            "SELECT CAST(current_setting('max_wal_senders') AS int8) AS max_wal_senders",
71            &[],
72        )
73        .await?;
74    Ok(max_wal_senders.get("max_wal_senders"))
75}
76
77pub async fn available_replication_slots(client: &Client) -> Result<i64, PostgresError> {
78    let available_replication_slots = client
79        .query_one(
80            "SELECT
81            CAST(current_setting('max_replication_slots') AS int8)
82              - (SELECT count(*) FROM pg_catalog.pg_replication_slots)
83              AS available_replication_slots;",
84            &[],
85        )
86        .await?;
87
88    let available_replication_slots: i64 =
89        available_replication_slots.get("available_replication_slots");
90
91    Ok(available_replication_slots)
92}
93
94/// Returns true if BYPASSRLS is set for the current user, false otherwise.
95///
96/// See <https://www.postgresql.org/docs/current/ddl-rowsecurity.html>
97pub async fn bypass_rls_attribute(client: &Client) -> Result<bool, PostgresError> {
98    let rls_attribute = client
99        .query_one(
100            "SELECT rolbypassrls FROM pg_roles WHERE rolname = CURRENT_USER;",
101            &[],
102        )
103        .await?;
104    Ok(rls_attribute.get("rolbypassrls"))
105}
106
107/// Returns an error if the tables identified by the oid's have RLS policies which
108/// affect the current user. Two checks are made:
109///
110/// 1. Identify which tables, from the provided oid's, have RLS policies that affecct the user or
111///    public.
112/// 2. If there are policies that affect the user, check if the BYPASSRLS attribute is set. If set,
113///    the role is unaffected by the policies.
114pub async fn validate_no_rls_policies(
115    client: &Client,
116    table_oids: &[Oid],
117) -> Result<(), PostgresError> {
118    if table_oids.is_empty() {
119        return Ok(());
120    }
121    let tables_with_rls_for_user = client
122        .query(
123            "SELECT
124                    format('%I.%I', pc.relnamespace::regnamespace, pc.relname) AS qualified_name
125                FROM pg_policy pp
126                JOIN pg_class pc ON pc.oid = polrelid
127                WHERE
128                    polrelid = ANY($1::oid[])
129                    AND
130                    (0 = ANY(polroles) OR CURRENT_USER::regrole::oid = ANY(polroles));",
131            &[&table_oids],
132        )
133        .await
134        .map_err(PostgresError::from)?;
135
136    let mut tables_with_rls_for_user = tables_with_rls_for_user
137        .into_iter()
138        .map(|row| row.get("qualified_name"))
139        .collect::<Vec<String>>();
140
141    // If the user has the BYPASSRLS flag set, then the policies don't apply, so we can
142    // return success.
143    if tables_with_rls_for_user.is_empty() || bypass_rls_attribute(client).await? {
144        Ok(())
145    } else {
146        tables_with_rls_for_user.sort();
147        Err(PostgresError::BypassRLSRequired(tables_with_rls_for_user))
148    }
149}
150
151pub async fn drop_replication_slots(
152    ssh_tunnel_manager: &SshTunnelManager,
153    config: Config,
154    slots: &[(&str, bool)],
155) -> Result<(), PostgresError> {
156    let client = config
157        .connect("postgres_drop_replication_slots", ssh_tunnel_manager)
158        .await?;
159    let replication_client = config.connect_replication(ssh_tunnel_manager).await?;
160    for (slot, should_wait) in slots {
161        let rows = client
162            .query(
163                "SELECT active_pid FROM pg_replication_slots WHERE slot_name = $1::TEXT",
164                &[&slot],
165            )
166            .await?;
167        match &*rows {
168            [] => {
169                // DROP_REPLICATION_SLOT will error if the slot does not exist
170                tracing::info!(
171                    "drop_replication_slots called on non-existent slot {}",
172                    slot
173                );
174                continue;
175            }
176            [row] => {
177                // The drop of a replication slot happens concurrently with an ingestion dataflow
178                // shutting down, therefore there is the possibility that the slot is still in use.
179                // We really don't want to leak the slot and not forcefully terminating the
180                // dataflow's connection risks timing out. For this reason we always kill the
181                // active backend and drop the slot.
182                let active_pid: Option<i32> = row.get("active_pid");
183                if let Some(active_pid) = active_pid {
184                    client
185                        .simple_query(&format!("SELECT pg_terminate_backend({active_pid})"))
186                        .await?;
187                }
188
189                let wait_str = if *should_wait { " WAIT" } else { "" };
190                let slot = Ident::new_unchecked(*slot).to_ast_string_simple();
191                replication_client
192                    .simple_query(&format!("DROP_REPLICATION_SLOT {slot}{wait_str}"))
193                    .await?;
194            }
195            _ => {
196                return Err(PostgresError::Generic(anyhow::anyhow!(
197                    "multiple pg_replication_slots entries for slot {}",
198                    &slot
199                )));
200            }
201        }
202    }
203    Ok(())
204}
205
206pub async fn get_timeline_id(client: &Client) -> Result<u64, PostgresError> {
207    if let Some(r) =
208        simple_query_opt(client, "SELECT timeline_id FROM pg_control_checkpoint()").await?
209    {
210        r.get("timeline_id")
211            .expect("Returns a row with a timeline ID")
212            .parse::<u64>()
213            .map_err(|err| {
214                PostgresError::Generic(anyhow::anyhow!(
215                    "Failed to parse timeline ID from IDENTIFY_SYSTEM: {}",
216                    err
217                ))
218            })
219    } else {
220        Err(PostgresError::Generic(anyhow::anyhow!(
221            "IDENTIFY_SYSTEM did not return a result row"
222        )))
223    }
224}
225
226pub async fn get_current_wal_lsn(client: &Client) -> Result<PgLsn, PostgresError> {
227    let row = client.query_one("SELECT pg_current_wal_lsn()", &[]).await?;
228    let lsn: PgLsn = row.get(0);
229
230    Ok(lsn)
231}
mz_postgres_util/replication.rs

mz_postgres_util/
replication.rs
