mz_postgres_util/

replication.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

use std::str::FromStr;

use tokio_postgres::{
    Client,
    types::{Oid, PgLsn},
};

use mz_ssh_util::tunnel_manager::SshTunnelManager;

use crate::{Config, PostgresError, Sql, query, query_one, simple_query, simple_query_opt};

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum WalLevel {
    Minimal,
    Replica,
    Logical,
}

impl std::str::FromStr for WalLevel {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "minimal" => Ok(Self::Minimal),
            "replica" => Ok(Self::Replica),
            "logical" => Ok(Self::Logical),
            o => Err(anyhow::anyhow!("unknown wal_level {}", o)),
        }
    }
}

impl std::fmt::Display for WalLevel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let s = match self {
            WalLevel::Minimal => "minimal",
            WalLevel::Replica => "replica",
            WalLevel::Logical => "logical",
        };

        f.write_str(s)
    }
}

#[mz_ore::test]
fn test_wal_level_max() {
    // Ensure `WalLevel::Logical` is the max among all levels.
    for o in [WalLevel::Minimal, WalLevel::Replica, WalLevel::Logical] {
        assert_eq!(WalLevel::Logical, WalLevel::Logical.max(o))
    }
}

pub async fn get_wal_level(client: &Client) -> Result<WalLevel, PostgresError> {
    let wal_level = query_one(client, crate::sql!("SHOW wal_level"), &[]).await?;
    let wal_level: String = wal_level.get("wal_level");
    Ok(WalLevel::from_str(&wal_level)?)
}

pub async fn get_max_wal_senders(client: &Client) -> Result<i64, PostgresError> {
    let max_wal_senders = query_one(
        client,
        crate::sql!("SELECT CAST(current_setting('max_wal_senders') AS int8) AS max_wal_senders"),
        &[],
    )
    .await?;
    Ok(max_wal_senders.get("max_wal_senders"))
}

pub async fn available_replication_slots(client: &Client) -> Result<i64, PostgresError> {
    let available_replication_slots = query_one(
        client,
        crate::sql!(
            "SELECT
            CAST(current_setting('max_replication_slots') AS int8)
              - (SELECT count(*) FROM pg_catalog.pg_replication_slots)
              AS available_replication_slots;"
        ),
        &[],
    )
    .await?;

    let available_replication_slots: i64 =
        available_replication_slots.get("available_replication_slots");

    Ok(available_replication_slots)
}

/// Returns true if BYPASSRLS is set for the current user, false otherwise.
///
/// See <https://www.postgresql.org/docs/current/ddl-rowsecurity.html>
pub async fn bypass_rls_attribute(client: &Client) -> Result<bool, PostgresError> {
    let rls_attribute = query_one(
        client,
        crate::sql!("SELECT rolbypassrls FROM pg_roles WHERE rolname = CURRENT_USER;"),
        &[],
    )
    .await?;
    Ok(rls_attribute.get("rolbypassrls"))
}

/// Returns an error if the tables identified by the oid's have RLS policies which
/// affect the current user. Two checks are made:
///
/// 1. Identify which tables, from the provided oid's, have RLS SELECT or ALL policies that affect
///    the user or public.
/// 2. If there are policies that affect the user, check if the BYPASSRLS attribute is set. If set,
///    the role is unaffected by the policies.
pub async fn validate_no_rls_policies(
    client: &Client,
    table_oids: &[Oid],
) -> Result<(), PostgresError> {
    if table_oids.is_empty() {
        return Ok(());
    }
    let tables_with_rls_for_user = query(
        client,
        crate::sql!(
            "SELECT
                    format('%I.%I', n.nspname, pc.relname) AS qualified_name
                FROM pg_policy pp
                JOIN pg_class pc ON pc.oid = pp.polrelid
                JOIN pg_namespace n ON n.oid = pc.relnamespace
                WHERE
                    pp.polrelid = ANY($1::oid[])
                    AND pp.polcmd IN ('r', '*')
                    AND (
                        0 = ANY(pp.polroles)
                        OR EXISTS (
                            SELECT 1
                            FROM unnest(pp.polroles) AS r(role_oid)
                            WHERE pg_has_role(CURRENT_USER, r.role_oid, 'USAGE')
                        )
                    );"
        ),
        &[&table_oids],
    )
    .await?;

    let mut tables_with_rls_for_user = tables_with_rls_for_user
        .into_iter()
        .map(|row| row.get("qualified_name"))
        .collect::<Vec<String>>();

    // If the user has the BYPASSRLS flag set, then the policies don't apply, so we can
    // return success.
    if tables_with_rls_for_user.is_empty() || bypass_rls_attribute(client).await? {
        Ok(())
    } else {
        tables_with_rls_for_user.sort();
        Err(PostgresError::BypassRLSRequired(tables_with_rls_for_user))
    }
}

pub async fn drop_replication_slots(
    ssh_tunnel_manager: &SshTunnelManager,
    config: Config,
    slots: &[(&str, bool)],
) -> Result<(), PostgresError> {
    let client = config
        .connect("postgres_drop_replication_slots", ssh_tunnel_manager)
        .await?;
    let replication_client = config.connect_replication(ssh_tunnel_manager).await?;
    for (slot, should_wait) in slots {
        let rows = query(
            &*client,
            crate::sql!("SELECT active_pid FROM pg_replication_slots WHERE slot_name = $1::TEXT"),
            &[&slot],
        )
        .await?;
        match &*rows {
            [] => {
                // DROP_REPLICATION_SLOT will error if the slot does not exist
                tracing::info!(
                    "drop_replication_slots called on non-existent slot {}",
                    slot
                );
                continue;
            }
            [row] => {
                // The drop of a replication slot happens concurrently with an ingestion dataflow
                // shutting down, therefore there is the possibility that the slot is still in use.
                // We really don't want to leak the slot and not forcefully terminating the
                // dataflow's connection risks timing out. For this reason we always kill the
                // active backend and drop the slot.
                let active_pid: Option<i32> = row.get("active_pid");
                if let Some(active_pid) = active_pid {
                    let query = crate::sql!("SELECT pg_terminate_backend({})", active_pid);
                    simple_query(&*client, query).await?;
                }

                let wait = if *should_wait {
                    crate::sql!(" WAIT")
                } else {
                    crate::sql!("")
                };
                let query = crate::sql!("DROP_REPLICATION_SLOT {}{}", Sql::ident(slot), wait);
                simple_query(&*replication_client, query).await?;
            }
            _ => {
                return Err(PostgresError::Generic(anyhow::anyhow!(
                    "multiple pg_replication_slots entries for slot {}",
                    &slot
                )));
            }
        }
    }
    Ok(())
}

pub async fn get_timeline_id(client: &Client) -> Result<u64, PostgresError> {
    if let Some(r) = simple_query_opt(
        client,
        crate::sql!("SELECT timeline_id FROM pg_control_checkpoint()"),
    )
    .await?
    {
        r.get("timeline_id")
            .expect("Returns a row with a timeline ID")
            .parse::<u64>()
            .map_err(|err| {
                PostgresError::Generic(anyhow::anyhow!(
                    "Failed to parse timeline ID from pg_control_checkpoint(): {}",
                    err
                ))
            })
    } else {
        Err(PostgresError::Generic(anyhow::anyhow!(
            "pg_control_checkpoint() did not return a result row"
        )))
    }
}

pub async fn fetch_max_lsn(
    client: &Client,
    is_physical_standby: bool,
) -> Result<PgLsn, PostgresError> {
    if is_physical_standby {
        // A physical standby will not support pg_current_wal_lsn, so we use pg_last_wal_replay_lsn. This reports
        // the latest LSN that the replica has successfully applied from the WAL.
        if let Some(lsn) = fetch_lsn(client, crate::sql!("SELECT pg_last_wal_replay_lsn()")).await?
        {
            return Ok(lsn);
        }
        return Err(PostgresError::Generic(anyhow::anyhow!(
            "pg_last_wal_replay_lsn unexpectedly None, expected client in recovery, is_in_recovery: {}",
            get_is_in_recovery(client).await?
        )));
    }
    // Based on the documentation, it appears that `pg_current_wal_lsn` has
    // the same "upper" semantics of `confirmed_flush_lsn`:
    // <https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-ADMIN-BACKUP>
    // We may need to revisit this and use `pg_current_wal_flush_lsn`.
    if let Some(lsn) = fetch_lsn(client, crate::sql!("SELECT pg_current_wal_lsn()")).await? {
        return Ok(lsn);
    }

    Err(PostgresError::Generic(anyhow::anyhow!(
        "pg_current_wal_lsn mysteriously has no value"
    )))
}

async fn fetch_lsn(client: &Client, query: Sql) -> Result<Option<PgLsn>, PostgresError> {
    let row = query_one(client, query, &[]).await?;
    let lsn: Option<PgLsn> = row.get(0);
    Ok(lsn)
}

/// Returns whether the server is in recovery, i.e. is a physical replica.
pub async fn get_is_in_recovery(client: &Client) -> Result<bool, PostgresError> {
    let row = query_one(client, crate::sql!("SELECT pg_is_in_recovery()"), &[]).await?;
    let is_in_recovery: Option<bool> = row.get(0);
    if let Some(is_in_recovery) = is_in_recovery {
        return Ok(is_in_recovery);
    }

    Err(PostgresError::Generic(anyhow::anyhow!(
        "pg_is_in_recovery() mysteriously has no value"
    )))
}