fix record of postgres version in databags

Author: BalabaDmitri

src/charm.py

@@ -71,6 +71,7 @@
     MONITORING_USER,
     PATRONI_CONF_PATH,
     PEER,
+    UPGRADE_RELATION,
     POSTGRESQL_SNAP_NAME,
     REPLICATION_PASSWORD_KEY,
     REWIND_PASSWORD_KEY,
@@ -95,7 +96,9 @@
 
 PRIMARY_NOT_REACHABLE_MESSAGE = "waiting for primary to be reachable from this unit"
 EXTENSIONS_DEPENDENCY_MESSAGE = "Unsatisfied plugin dependencies. Please check the logs"
-DIFFERENT_VERSIONS_PSQL_BLOCKING_MESSAGE = "Please select the correct version of postgresql to use. No need to use different versions of postgresql."
+DIFFERENT_VERSIONS_PSQL_BLOCKING_MESSAGE = (
+    "Please select the correct version of postgresql to use. You cannot use different versions of postgresql!"
+)
 
 Scopes = Literal[APP_SCOPE, UNIT_SCOPE]
 
@@ -157,6 +160,7 @@ def __init__(self, *args):
         self.framework.observe(self.on.config_changed, self._on_config_changed)
         self.framework.observe(self.on.get_primary_action, self._on_get_primary)
         self.framework.observe(self.on[PEER].relation_changed, self._on_peer_relation_changed)
+        self.framework.observe(self.on[UPGRADE_RELATION].relation_changed, self._on_upgrade_relation_changed)
         self.framework.observe(self.on.secret_changed, self._on_peer_relation_changed)
         self.framework.observe(self.on.secret_remove, self._on_peer_relation_changed)
         self.framework.observe(self.on[PEER].relation_departed, self._on_peer_relation_departed)
@@ -535,6 +539,17 @@ def _on_peer_relation_changed(self, event: HookEvent):
 
         self._validate_database_version()
 
+    def _on_upgrade_relation_changed(self, event: HookEvent):
+        if not self.unit.is_leader():
+            return
+
+        if self.upgrade.idle:
+            logger.debug("Defer _on_upgrade_relation_changed: upgrade in progress")
+            event.defer()
+            return
+
+        self._set_workload_version(self._patroni.get_postgresql_version())
+
     # Split off into separate function, because of complexity _on_peer_relation_changed
     def _start_stop_pgbackrest_service(self, event: HookEvent) -> None:
         # Start or stop the pgBackRest TLS server service when TLS certificate change.
@@ -1020,7 +1035,7 @@ def _on_start(self, event: StartEvent) -> None:
 
         self.unit_peer_data.update({"ip": self.get_hostname_by_unit(None)})
 
-        self.unit.set_workload_version(self._patroni.get_postgresql_version())
+        self._set_workload_version(self._patroni.get_postgresql_version())
 
         # Open port
         try:
@@ -1596,6 +1611,12 @@ def client_relations(self) -> List[Relation]:
                 relations.append(relation)
         return relations
 
+    def _set_workload_version(self, psql_version):
+        """Record the version of the software running as the workload. Also writes the version into the databags"""
+        self.unit.set_workload_version(psql_version)
+        if self.unit.is_leader():
+            self.app_peer_data.update({"database-version": psql_version})
+
     def _validate_database_version(self):
         """Checking that only one version of Postgres is used."""
         peer_db_version = self.app_peer_data.get("database-version")

src/constants.py

@@ -10,6 +10,7 @@
 LEGACY_DB = "db"
 LEGACY_DB_ADMIN = "db-admin"
 PEER = "database-peers"
+UPGRADE_RELATION = "upgrade"
 ALL_CLIENT_RELATIONS = [DATABASE, LEGACY_DB, LEGACY_DB_ADMIN]
 API_REQUEST_TIMEOUT = 5
 PATRONI_CLUSTER_STATUS_ENDPOINT = "cluster"

tests/integration/ha_tests/helpers.py

@@ -38,6 +38,7 @@
 PATRONI_SERVICE_DEFAULT_PATH = f"/etc/systemd/system/{SERVICE_NAME}"
 RESTART_CONDITION = "no"
 ORIGINAL_RESTART_CONDITION = "always"
+SECOND_APPLICATION = "second-cluster"
 
 
 class MemberNotListedOnClusterError(Exception):
@@ -872,10 +873,21 @@ async def reused_full_cluster_recovery_storage(ops_test: OpsTest, unit_name) ->
 
 
 async def get_db_connection(ops_test, dbname, is_primary=True, replica_unit_name=""):
+    """Returns a PostgreSQL connection string.
+
+       Args:
+           ops_test: The ops test framework instance
+           dbname: The name of the database
+           is_primary: Whether to use a primary unit (default is True, so it uses the primary
+           replica_unit_name: The name of the replica unit
+
+       Returns:
+           a PostgreSQL connection string
+       """
     unit_name = await get_primary(ops_test, APP_NAME)
     password = await get_password(ops_test, APP_NAME)
     address = get_unit_address(ops_test, unit_name)
-    if not is_primary and unit_name != "":
+    if not is_primary and replica_unit_name != "":
         unit_name = replica_unit_name
         address = ops_test.model.applications[APP_NAME].units[unit_name].public_address
     connection_string = (
@@ -886,6 +898,11 @@ async def get_db_connection(ops_test, dbname, is_primary=True, replica_unit_name
 
 
 async def validate_test_data(connection_string):
+    """Checking test data.
+
+        Args:
+          connection_string: Database connection string
+        """
     with psycopg2.connect(connection_string) as connection:
         connection.autocommit = True
         with connection.cursor() as cursor:
@@ -896,6 +913,11 @@ async def validate_test_data(connection_string):
 
 
 async def create_test_data(connection_string):
+    """Creating test data in the database.
+
+        Args:
+          connection_string: Database connection string
+        """
     with psycopg2.connect(connection_string) as connection:
         connection.autocommit = True
         with connection.cursor() as cursor:
@@ -911,6 +933,16 @@ async def create_test_data(connection_string):
 
 
 async def get_last_added_unit(ops_test, app, prev_units):
+    """Returns a unit.
+
+        Args:
+          ops_test: The ops test framework instance
+          app: The name of the application
+          prev_units: List of unit names before adding the last unit
+
+        Returns:
+          last added unit
+        """
     curr_units = [unit.name for unit in ops_test.model.applications[app].units]
     new_unit = list(set(curr_units) - set(prev_units))[0]
     for unit in ops_test.model.applications[app].units:

tests/integration/ha_tests/test_restore_cluster.py

@@ -20,10 +20,10 @@
     add_unit_with_storage,
     reused_full_cluster_recovery_storage,
     storage_id,
+    SECOND_APPLICATION,
 )
 
 FIRST_APPLICATION = "first-cluster"
-SECOND_APPLICATION = "second-cluster"
 
 logger = logging.getLogger(__name__)
 

tests/integration/ha_tests/test_self_healing.py

@@ -56,8 +56,8 @@
     update_restart_condition,
     validate_test_data,
     wait_network_restore,
+    SECOND_APPLICATION,
 )
-from .test_restore_cluster import SECOND_APPLICATION
 
 logger = logging.getLogger(__name__)
 

tests/integration/helpers.py

@@ -35,7 +35,6 @@
 DATABASE_APP_NAME = METADATA["name"]
 STORAGE_PATH = METADATA["storage"]["pgdata"]["location"]
 APPLICATION_NAME = "postgresql-test-app"
-FIRST_DATABASE_RELATION_NAME = "first-database"
 
 
 async def build_connection_string(