Logical Replication

actions.yaml

@@ -1,6 +1,16 @@
 # Copyright 2021 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.
   params:
@@ -30,6 +40,8 @@ get-password:
         Possible values - operator, replication, rewind, patroni.
 list-backups:
   description: Lists backups in s3 storage.
+list-publications:
+list-subscriptions:
 pre-upgrade-check:
   description: Run necessary pre-upgrade checks and preparations before executing a charm refresh.
 promote-to-primary:
@@ -43,6 +55,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.
@@ -70,3 +86,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 = 51
+LIBPATCH = 52
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -124,6 +124,33 @@ class PostgreSQLListUsersError(Exception):
 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."""
@@ -778,6 +805,134 @@ 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_group_map(group_map: Optional[str]) -> List[Tuple]:
         """Build the PostgreSQL authorization group-map.

metadata.yaml

@@ -30,6 +30,9 @@ provides:
     interface: postgresql_async
     limit: 1
     optional: true
+  logical-replication-offer:
+    interface: postgresql_logical_replication
+    optional: true
   database:
     interface: postgresql_client
   db:
@@ -45,6 +48,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

@@ -44,6 +44,11 @@
     PGBACKREST_LOGS_PATH,
     POSTGRESQL_DATA_PATH,
 )
+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__)
 
@@ -1219,13 +1224,35 @@ def _pre_restore_checks(self, event: ActionEvent) -> bool:
             event.fail(error_message)
             return False
 
+        logger.info("Checking that the cluster is not replicating data to a standby cluster")
+        for relation in [
+            self.model.get_relation(REPLICATION_CONSUMER_RELATION),
+            self.model.get_relation(REPLICATION_OFFER_RELATION),
+        ]:
+            if not relation:
+                continue
+            error_message = "Unit cannot restore backup as the cluster is replicating data to a standby cluster"
+            logger.error(f"Restore failed: {error_message}")
+            event.fail(error_message)
+            return False
+
         logger.info("Checking that this unit was already elected the leader unit")
         if not self.charm.unit.is_leader():
             error_message = "Unit cannot restore backup as it was not elected the leader unit yet"
             logger.error(f"Restore failed: {error_message}")
             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

@@ -110,6 +110,7 @@
     PostgreSQLAsyncReplication,
 )
 from relations.db import EXTENSIONS_BLOCKING_MESSAGE, DbProvides
+from relations.logical_replication import PostgreSQLLogicalReplication
 from relations.postgresql_provider import PostgreSQLProvider
 from rotate_logs import RotateLogs
 from upgrade import PostgreSQLUpgrade, get_postgresql_dependencies_model
@@ -214,6 +215,7 @@ def __init__(self, *args):
         self.ldap = PostgreSQLLDAP(self, "ldap")
         self.tls = PostgreSQLTLS(self, PEER)
         self.async_replication = PostgreSQLAsyncReplication(self)
+        self.logical_replication = PostgreSQLLogicalReplication(self)
         self.restart_manager = RollingOpsManager(
             charm=self, relation="restart", callback=self._restart
         )
@@ -1997,6 +1999,12 @@ def update_config(self, is_creating_backup: bool = False, no_peers: bool = False
             self.model.config, self.get_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
+        )
+
         # Update and reload configuration based on TLS files availability.
         self._patroni.render_patroni_yml_file(
             connectivity=self.is_connectivity_enabled,
@@ -2011,6 +2019,7 @@ def update_config(self, is_creating_backup: bool = False, no_peers: bool = False
             restore_stanza=self.app_peer_data.get("restore-stanza"),
             parameters=pg_parameters,
             no_peers=no_peers,
+            slots=replication_slots_json,
         )
         if no_peers:
             return True
@@ -2047,6 +2056,8 @@ def update_config(self, is_creating_backup: bool = False, no_peers: bool = False
             "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(enable_tls)
 
         cache = snap.SnapCache()