Logical Replication

actions.yaml

@@ -1,6 +1,16 @@
 # Copyright 2023 Canonical Ltd.
 # See LICENSE file for licensing details.
 
+# TODO: add descriptions for logical replication actions
+
+add-publication:
+  params:
+    name:
+      type: string
+    database:
+      type: string
+    tables:
+      type: string
 create-backup:
   description: Creates a backup to s3 storage in AWS.
   params:
@@ -31,6 +41,8 @@ get-password:
         Possible values - backup, operator, replication, rewind, patroni.
 list-backups:
   description: Lists backups in s3 storage in AWS.
+list-publications:
+list-subscriptions:
 pre-upgrade-check:
   description: Run necessary pre-upgrade checks and preparations before executing a charm refresh.
 promote-to-primary:
@@ -44,6 +56,10 @@ promote-to-primary:
     force:
       type: boolean
       description: Force the promotion of a cluster when there is already a primary cluster.
+remove-publication:
+  params:
+    name:
+      type: string
 restore:
   description: Restore a database backup using pgBackRest.
     S3 credentials are retrieved from a relation with the S3 integrator charm.
@@ -73,3 +89,11 @@ set-tls-private-key:
     private-key:
       type: string
       description: The content of private key for communications with clients. Content will be auto-generated if this option is not specified.
+subscribe:
+  params:
+    name:
+      type: string
+unsubscribe:
+  params:
+    name:
+      type: string

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -35,7 +35,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 47
+LIBPATCH = 48
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -125,6 +125,42 @@ class PostgreSQLUpdateUserPasswordError(Exception):
     """Exception raised when updating a user password fails."""
 
 
+class PostgreSQLDatabaseExistsError(Exception):
+    """Exception raised during database existence check."""
+
+
+class PostgreSQLTableExistsError(Exception):
+    """Exception raised during table existence check."""
+
+
+class PostgreSQLIsTableEmptyError(Exception):
+    """Exception raised during table emptiness check."""
+
+
+class PostgreSQLCreatePublicationError(Exception):
+    """Exception raised when creating PostgreSQL publication."""
+
+
+class PostgreSQLDropPublicationError(Exception):
+    """Exception raised when dropping PostgreSQL publication."""
+
+
+class PostgreSQLCreateSubscriptionError(Exception):
+    """Exception raised when creating PostgreSQL subscription."""
+
+
+class PostgreSQLSubscriptionExistsError(Exception):
+    """Exception raised during subscription existence check."""
+
+
+class PostgreSQLUpdateSubscriptionError(Exception):
+    """Exception raised when updating PostgreSQL subscription."""
+
+
+class PostgreSQLDropSubscriptionError(Exception):
+    """Exception raised when dropping PostgreSQL subscription."""
+
+
 class PostgreSQL:
     """Class to encapsulate all operations related to interacting with PostgreSQL instance."""
 
@@ -773,6 +809,165 @@ def is_restart_pending(self) -> bool:
             if connection:
                 connection.close()
 
+    def database_exists(self, db: str) -> bool:
+        """Check whether specified database exists."""
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL("SELECT datname FROM pg_database WHERE datname={};").format(Literal(db))
+                )
+                return cursor.fetchone() is not None
+        except psycopg2.Error as e:
+            logger.error(f"Failed to check Postgresql database existence: {e}")
+            raise PostgreSQLDatabaseExistsError() from e
+
+    def table_exists(self, db: str, schema: str, table: str) -> bool:
+        """Check whether specified table in database exists."""
+        try:
+            with self._connect_to_database(
+                database=db
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL(
+                        "SELECT tablename FROM pg_tables WHERE schemaname={} AND tablename={};"
+                    ).format(Literal(schema), Literal(table))
+                )
+                return cursor.fetchone() is not None
+        except psycopg2.Error as e:
+            logger.error(f"Failed to check Postgresql table existence: {e}")
+            raise PostgreSQLTableExistsError() from e
+
+    def is_table_empty(self, db: str, schema: str, table: str) -> bool:
+        """Check whether table is empty."""
+        try:
+            with self._connect_to_database(
+                database=db
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(SQL("SELECT COUNT(1) FROM {};").format(Identifier(schema, table)))
+                return cursor.fetchone()[0] == 0
+        except psycopg2.Error as e:
+            logger.error(f"Failed to check whether table is empty: {e}")
+            raise PostgreSQLIsTableEmptyError() from e
+
+    def create_publication(self, db: str, name: str, schematables: list[str]) -> None:
+        """Create PostgreSQL publication."""
+        try:
+            with self._connect_to_database(
+                database=db
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL("CREATE PUBLICATION {} FOR TABLE {};").format(
+                        Identifier(name),
+                        SQL(",").join(
+                            Identifier(schematable.split(".")[0], schematable.split(".")[1])
+                            for schematable in schematables
+                        ),
+                    )
+                )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to create Postgresql publication: {e}")
+            raise PostgreSQLCreatePublicationError() from e
+
+    def drop_publication(self, db: str, publication: str) -> None:
+        """Drop PostgreSQL publication."""
+        try:
+            with self._connect_to_database(
+                database=db
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL("DROP PUBLICATION IF EXISTS {};").format(
+                        Identifier(publication),
+                    )
+                )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to drop Postgresql publication: {e}")
+            raise PostgreSQLDropPublicationError() from e
+
+    def create_subscription(
+        self,
+        subscription: str,
+        host: str,
+        db: str,
+        user: str,
+        password: str,
+        replication_slot: str,
+    ) -> None:
+        """Create PostgreSQL subscription."""
+        try:
+            with self._connect_to_database(
+                database=db
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL(
+                        "CREATE SUBSCRIPTION {} CONNECTION {} PUBLICATION {} WITH (copy_data=true,create_slot=false,enabled=true,slot_name={});"
+                    ).format(
+                        Identifier(subscription),
+                        Literal(f"host={host} dbname={db} user={user} password={password}"),
+                        Identifier(subscription),
+                        Identifier(replication_slot),
+                    )
+                )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to create Postgresql subscription: {e}")
+            raise PostgreSQLCreateSubscriptionError() from e
+
+    def subscription_exists(self, db: str, subscription: str) -> bool:
+        """Check whether specified subscription in database exists."""
+        try:
+            with self._connect_to_database(
+                database=db
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL("SELECT subname FROM pg_subscription WHERE subname={};").format(
+                        Literal(subscription)
+                    )
+                )
+                return cursor.fetchone() is not None
+        except psycopg2.Error as e:
+            logger.error(f"Failed to check Postgresql subscription existence: {e}")
+            raise PostgreSQLSubscriptionExistsError() from e
+
+    def update_subscription(self, db: str, subscription: str, host: str, user: str, password: str):
+        """Update PostgreSQL subscription connection details."""
+        try:
+            with self._connect_to_database(
+                database=db
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL("ALTER SUBSCRIPTION {} CONNECTION {}").format(
+                        Identifier(subscription),
+                        Literal(f"host={host} dbname={db} user={user} password={password}"),
+                    )
+                )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to update Postgresql subscription: {e}")
+            raise PostgreSQLUpdateSubscriptionError() from e
+
+    def drop_subscription(self, db: str, subscription: str) -> None:
+        """Drop PostgreSQL subscription."""
+        try:
+            with self._connect_to_database(
+                database=db
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL("ALTER SUBSCRIPTION {} DISABLE;").format(
+                        Identifier(subscription),
+                    )
+                )
+                cursor.execute(
+                    SQL("ALTER SUBSCRIPTION {} SET (slot_name=NONE);").format(
+                        Identifier(subscription),
+                    )
+                )
+                cursor.execute(
+                    SQL("DROP SUBSCRIPTION {};").format(
+                        Identifier(subscription),
+                    )
+                )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to drop Postgresql subscription: {e}")
+            raise PostgreSQLDropSubscriptionError() from e
+
     @staticmethod
     def build_postgresql_parameters(
         config_options: dict, available_memory: int, limit_memory: Optional[int] = None

metadata.yaml

@@ -43,6 +43,9 @@ provides:
     interface: postgresql_async
     limit: 1
     optional: true
+  logical-replication-offer:
+    interface: postgresql_logical_replication
+    optional: true
   database:
     interface: postgresql_client
   db:
@@ -59,6 +62,10 @@ requires:
     interface: postgresql_async
     limit: 1
     optional: true
+  logical-replication:
+    interface: postgresql_logical_replication
+    limit: 1
+    optional: true
   certificates:
     interface: tls-certificates
     limit: 1

src/backups.py

@@ -35,6 +35,10 @@
     WORKLOAD_OS_USER,
 )
 from relations.async_replication import REPLICATION_CONSUMER_RELATION, REPLICATION_OFFER_RELATION
+from relations.logical_replication import (
+    LOGICAL_REPLICATION_OFFER_RELATION,
+    LOGICAL_REPLICATION_RELATION,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -1082,7 +1086,7 @@ def _fetch_backup_from_id(self, backup_id: str) -> str:
 
         return None
 
-    def _pre_restore_checks(self, event: ActionEvent) -> bool:
+    def _pre_restore_checks(self, event: ActionEvent) -> bool:  # noqa: C901
         """Run some checks before starting the restore.
 
         Returns:
@@ -1156,6 +1160,16 @@ def _pre_restore_checks(self, event: ActionEvent) -> bool:
             event.fail(error_message)
             return False
 
+        if self.model.get_relation(LOGICAL_REPLICATION_RELATION) or len(
+            self.model.relations.get(LOGICAL_REPLICATION_OFFER_RELATION, ())
+        ):
+            error_message = (
+                "Cannot proceed with restore with an active logical replication connection"
+            )
+            logger.error(f"Restore failed: {error_message}")
+            event.fail(error_message)
+            return False
+
         return True
 
     def _render_pgbackrest_conf_file(self) -> bool:

src/charm.py

@@ -15,6 +15,8 @@
 from pathlib import Path
 from typing import Literal, get_args
 
+from relations.logical_replication import PostgreSQLLogicalReplication
+
 # First platform-specific import, will fail on wrong architecture
 try:
     import psycopg2
@@ -237,6 +239,7 @@ def __init__(self, *args):
         self.ldap = PostgreSQLLDAP(self, "ldap")
         self.tls = PostgreSQLTLS(self, PEER, [self.primary_endpoint, self.replicas_endpoint])
         self.async_replication = PostgreSQLAsyncReplication(self)
+        self.logical_replication = PostgreSQLLogicalReplication(self)
         self.restart_manager = RollingOpsManager(
             charm=self, relation="restart", callback=self._restart
         )
@@ -1962,6 +1965,12 @@ def update_config(self, is_creating_backup: bool = False) -> bool:
             self.model.config, available_memory, limit_memory
         )
 
+        replication_slots_json = (
+            json.loads(self.app_peer_data["replication-slots"])
+            if "replication-slots" in self.app_peer_data
+            else None
+        )
+
         logger.info("Updating Patroni config file")
         # Update and reload configuration based on TLS files availability.
         self._patroni.render_patroni_yml_file(
@@ -1977,6 +1986,7 @@ def update_config(self, is_creating_backup: bool = False) -> bool:
             stanza=self.app_peer_data.get("stanza", self.unit_peer_data.get("stanza")),
             restore_stanza=self.app_peer_data.get("restore-stanza"),
             parameters=postgresql_parameters,
+            slots=replication_slots_json,
         )
 
         if not self._is_workload_running:
@@ -2008,6 +2018,8 @@ def update_config(self, is_creating_backup: bool = False) -> bool:
             "wal_keep_size": self.config.durability_wal_keep_size,
         })
 
+        self._patroni.ensure_slots_controller_by_patroni(replication_slots_json or {})
+
         self._handle_postgresql_restart_need()
 
         # Restart the monitoring service if the password was rotated