database_replicator/postgres/
privileges.rs

1// ABOUTME: Privilege checking utilities for migration prerequisites
2// ABOUTME: Validates source and target databases have required permissions
3
4use anyhow::{Context, Result};
5use tokio_postgres::Client;
6
7/// Result of privilege check for a PostgreSQL user
8///
9/// Contains information about the user's permissions required for migration.
10pub struct PrivilegeCheck {
11    /// User has REPLICATION privilege (required for source database)
12    pub has_replication: bool,
13    /// User has CREATEDB privilege (required for target database)
14    pub has_create_db: bool,
15    /// User has CREATEROLE privilege (optional, for role migration)
16    pub has_create_role: bool,
17    /// User is a superuser (bypasses other privilege requirements)
18    pub is_superuser: bool,
19}
20
21/// Check if connected user has replication privileges (needed for source)
22///
23/// Queries `pg_roles` to determine the privileges of the currently connected user.
24/// For source databases, the user must have REPLICATION privilege (or be a superuser)
25/// to enable logical replication.
26///
27/// # Arguments
28///
29/// * `client` - Connected PostgreSQL client
30///
31/// # Returns
32///
33/// Returns a `PrivilegeCheck` containing the user's privileges.
34///
35/// # Errors
36///
37/// This function will return an error if the database query fails.
38///
39/// # Examples
40///
41/// ```no_run
42/// # use anyhow::Result;
43/// # use database_replicator::postgres::{connect, check_source_privileges};
44/// # async fn example() -> Result<()> {
45/// let client = connect("postgresql://user:pass@localhost:5432/mydb").await?;
46/// let privs = check_source_privileges(&client).await?;
47/// assert!(privs.has_replication || privs.is_superuser);
48/// # Ok(())
49/// # }
50/// ```
51pub async fn check_source_privileges(client: &Client) -> Result<PrivilegeCheck> {
52    let row = client
53        .query_one(
54            "SELECT rolreplication, rolcreatedb, rolcreaterole, rolsuper
55             FROM pg_roles
56             WHERE rolname = current_user",
57            &[],
58        )
59        .await
60        .context("Failed to query user privileges")?;
61
62    Ok(PrivilegeCheck {
63        has_replication: row.get(0),
64        has_create_db: row.get(1),
65        has_create_role: row.get(2),
66        is_superuser: row.get(3),
67    })
68}
69
70/// Check if connected user has sufficient privileges for target database
71///
72/// Queries `pg_roles` to determine the privileges of the currently connected user.
73/// For target databases, the user must have CREATEDB privilege (or be a superuser)
74/// to create new databases during migration.
75///
76/// # Arguments
77///
78/// * `client` - Connected PostgreSQL client
79///
80/// # Returns
81///
82/// Returns a `PrivilegeCheck` containing the user's privileges.
83///
84/// # Errors
85///
86/// This function will return an error if the database query fails.
87///
88/// # Examples
89///
90/// ```no_run
91/// # use anyhow::Result;
92/// # use database_replicator::postgres::{connect, check_target_privileges};
93/// # async fn example() -> Result<()> {
94/// let client = connect("postgresql://user:pass@localhost:5432/mydb").await?;
95/// let privs = check_target_privileges(&client).await?;
96/// assert!(privs.has_create_db || privs.is_superuser);
97/// # Ok(())
98/// # }
99/// ```
100pub async fn check_target_privileges(client: &Client) -> Result<PrivilegeCheck> {
101    // Same query as source
102    check_source_privileges(client).await
103}
104
105/// Check the wal_level setting on the target database
106///
107/// Queries the current `wal_level` configuration parameter.
108/// For logical replication (subscriptions), `wal_level` must be set to `logical`.
109///
110/// # Arguments
111///
112/// * `client` - Connected PostgreSQL client
113///
114/// # Returns
115///
116/// Returns the current `wal_level` setting as a String (e.g., "replica", "logical").
117///
118/// # Errors
119///
120/// This function will return an error if the database query fails.
121///
122/// # Examples
123///
124/// ```no_run
125/// # use anyhow::Result;
126/// # use database_replicator::postgres::{connect, check_wal_level};
127/// # async fn example() -> Result<()> {
128/// let client = connect("postgresql://user:pass@localhost:5432/mydb").await?;
129/// let wal_level = check_wal_level(&client).await?;
130/// assert_eq!(wal_level, "logical");
131/// # Ok(())
132/// # }
133/// ```
134pub async fn check_wal_level(client: &Client) -> Result<String> {
135    let row = client
136        .query_one("SHOW wal_level", &[])
137        .await
138        .context("Failed to query wal_level setting")?;
139
140    let wal_level: String = row.get(0);
141    Ok(wal_level)
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use crate::postgres::connect;
148
149    #[tokio::test]
150    #[ignore]
151    async fn test_check_source_privileges() {
152        let url = std::env::var("TEST_SOURCE_URL").unwrap();
153        let client = connect(&url).await.unwrap();
154
155        let privileges = check_source_privileges(&client).await.unwrap();
156
157        // Should have at least one privilege
158        assert!(
159            privileges.has_replication || privileges.is_superuser,
160            "Source user should have REPLICATION privilege or be superuser"
161        );
162    }
163
164    #[tokio::test]
165    #[ignore]
166    async fn test_check_target_privileges() {
167        let url = std::env::var("TEST_TARGET_URL").unwrap();
168        let client = connect(&url).await.unwrap();
169
170        let privileges = check_target_privileges(&client).await.unwrap();
171
172        // Should have create privileges or be superuser
173        assert!(
174            privileges.has_create_db || privileges.is_superuser,
175            "Target user should have CREATE DATABASE privilege or be superuser"
176        );
177    }
178}