Add password rotation action

Author: marceloneppel

actions.yaml

@@ -6,3 +6,13 @@ get-primary:
 get-operator-password:
   description: Get the operator user password used by charm.
     It is internal charm user, SHOULD NOT be used by applications.
+rotate-users-passwords:
+  description: Change the operator user password used by charm.
+    It is internal charm user, SHOULD NOT be used by applications.
+  params:
+    user:
+      type: string
+      description: The password will be auto-generated if this option is not specified.
+    password:
+      type: string
+      description: The password will be auto-generated if this option is not specified.

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -19,6 +19,7 @@
 Any charm using this library should import the `psycopg2` or `psycopg2-binary` dependency.
 """
 import logging
+from typing import Dict
 
 import psycopg2
 from psycopg2 import sql
@@ -31,7 +32,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 1
+LIBPATCH = 2
 
 
 logger = logging.getLogger(__name__)
@@ -53,6 +54,10 @@ class PostgreSQLGetPostgreSQLVersionError(Exception):
     """Exception raised when retrieving PostgreSQL version fails."""
 
 
+class PostgreSQLSetUserPasswordError(Exception):
+    """Exception raised when setting/updating a user password fails."""
+
+
 class PostgreSQL:
     """Class to encapsulate all operations related to interacting with PostgreSQL instance."""
 
@@ -68,20 +73,25 @@ def __init__(
         self.password = password
         self.database = database
 
-    def _connect_to_database(self, database: str = None) -> psycopg2.extensions.connection:
+    def _connect_to_database(
+        self, database: str = None, autocommit: bool = True
+    ) -> psycopg2.extensions.connection:
         """Creates a connection to the database.
 
         Args:
             database: database to connect to (defaults to the database
                 provided when the object for this class was created).
+            autocommit: whether every command issued to the database
+                should be immediately committed.
 
         Returns:
              psycopg2 connection object.
         """
         connection = psycopg2.connect(
             f"dbname='{database if database else self.database}' user='{self.user}' host='{self.host}' password='{self.password}' connect_timeout=1"
         )
-        connection.autocommit = True
+        if autocommit:
+            connection.autocommit = True
         return connection
 
     def create_database(self, database: str, user: str) -> None:
@@ -173,3 +183,30 @@ def get_postgresql_version(self) -> str:
         except psycopg2.Error as e:
             logger.error(f"Failed to get PostgreSQL version: {e}")
             raise PostgreSQLGetPostgreSQLVersionError()
+
+    def rotate_users_passwords(self, users_with_passwords: Dict[str, str]) -> None:
+        """Rotates one or more users passwords.
+
+        Args:
+            users_with_passwords: a dict following the format {"user": "password"}.
+                It can contain multiple users.
+
+        Raises:
+            PostgreSQLSetUserPasswordError if any user password couldn't be changed.
+        """
+        try:
+            with self._connect_to_database(
+                autocommit=False
+            ) as connection, connection.cursor() as cursor:
+                for user, password in users_with_passwords.items():
+                    cursor.execute(
+                        sql.SQL(
+                            "ALTER USER {} WITH ENCRYPTED PASSWORD '" + password + "';"
+                        ).format(sql.Identifier(user))
+                    )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to rotate user password: {e}")
+            raise PostgreSQLSetUserPasswordError()
+        finally:
+            if connection is not None:
+                connection.close()

src/charm.py

@@ -7,7 +7,10 @@
 import logging
 from typing import List
 
-from charms.postgresql_k8s.v0.postgresql import PostgreSQL
+from charms.postgresql_k8s.v0.postgresql import (
+    PostgreSQL,
+    PostgreSQLSetUserPasswordError,
+)
 from lightkube import ApiError, Client, codecs
 from lightkube.resources.core_v1 import Pod
 from ops.charm import (
@@ -30,7 +33,7 @@
 from requests import ConnectionError
 from tenacity import RetryError
 
-from constants import PEER, USER
+from constants import PEER, REPLICATION_USER, SYSTEM_USERS, USER
 from patroni import NotReadyError, Patroni
 from relations.db import DbProvides
 from relations.postgresql_provider import PostgreSQLProvider
@@ -61,6 +64,9 @@ def __init__(self, *args):
         self.framework.observe(
             self.on.get_operator_password_action, self._on_get_operator_password
         )
+        self.framework.observe(
+            self.on.rotate_users_passwords_action, self._on_rotate_users_passwords
+        )
         self.framework.observe(self.on.get_primary_action, self._on_get_primary)
         self.framework.observe(self.on.update_status, self._on_update_status)
         self._storage_path = self.meta.storages["pgdata"].location
@@ -391,6 +397,53 @@ def _on_get_operator_password(self, event: ActionEvent) -> None:
         """Returns the password for the operator user as an action response."""
         event.set_results({"operator-password": self._get_operator_password()})
 
+    def _on_rotate_users_passwords(self, event: ActionEvent) -> None:
+        """Rotate the password for all system users or the specified user."""
+        # Only the leader can write the new password into peer relation.
+        if not self.unit.is_leader():
+            event.fail("The action can be run only on the leader unit")
+            return
+
+        if "user" in event.params:
+            user = event.params["user"]
+
+            # Fail if the user is not a system user.
+            # One example is users created through relations.
+            if user not in SYSTEM_USERS:
+                event.fail(f"User {user} is not a system user")
+                return
+
+            # Generate a new password and use it if no password was provided to the action.
+            users = {
+                user: event.params["password"] if "password" in event.params else new_password()
+            }
+        else:
+            if "password" in event.params:
+                event.fail("The same password cannot be set for multiple users")
+                return
+
+            users = {user: new_password() for user in SYSTEM_USERS}
+
+        try:
+            self.postgresql.rotate_users_passwords(users)
+        except PostgreSQLSetUserPasswordError as e:
+            event.fail(f"Failed to set user password with error {e}")
+            return
+
+        # Update the password in the peer relation if the operation was successful.
+        for user, password in users.items():
+            self._peers.data[self.app].update({f"{user}-password": password})
+
+        # for unit in
+        self._patroni.reload_patroni_configuration()
+
+        # Return the generated password when the user option is given.
+        if "user" in event.params and "password" not in event.params:
+            user = event.params["user"]
+            event.set_results(
+                {f"{user}-password": self._peers.data[self.app].get(f"{user}-password")}
+            )
+
     def _on_get_primary(self, event: ActionEvent) -> None:
         """Get primary instance."""
         try:
@@ -499,7 +552,7 @@ def _postgresql_layer(self) -> Layer:
                         "PATRONI_KUBERNETES_USE_ENDPOINTS": "true",
                         "PATRONI_NAME": pod_name,
                         "PATRONI_SCOPE": f"patroni-{self._name}",
-                        "PATRONI_REPLICATION_USERNAME": "replication",
+                        "PATRONI_REPLICATION_USERNAME": REPLICATION_USER,
                         "PATRONI_REPLICATION_PASSWORD": self._replication_password,
                         "PATRONI_SUPERUSER_USERNAME": USER,
                         "PATRONI_SUPERUSER_PASSWORD": self._get_operator_password(),

src/constants.py

@@ -5,4 +5,6 @@
 
 DATABASE_PORT = "5432"
 PEER = "database-peers"
+REPLICATION_USER = "replication"
 USER = "operator"
+SYSTEM_USERS = [REPLICATION_USER, USER]