[DPE-4106] Tests legacy and modern endpoints simultaneously

src/constants.py

@@ -11,6 +11,7 @@
 LEGACY_DB_ADMIN = "db-admin"
 PEER = "database-peers"
 ALL_CLIENT_RELATIONS = [DATABASE, LEGACY_DB, LEGACY_DB_ADMIN]
+ALL_LEGACY_RELATIONS = [LEGACY_DB, LEGACY_DB_ADMIN]
 API_REQUEST_TIMEOUT = 5
 PATRONI_CLUSTER_STATUS_ENDPOINT = "cluster"
 BACKUP_USER = "backup"
@@ -67,3 +68,7 @@
 UNIT_SCOPE = "unit"
 
 SECRET_KEY_OVERRIDES = {"ca": "cauth"}
+
+ENDPOINT_SIMULTANEOUSLY_BLOCKING_MESSAGE = (
+    "Please choose one endpoint to use. No need to relate all of them simultaneously!"
+)

src/relations/db.py

@@ -21,7 +21,12 @@
 from ops.model import ActiveStatus, BlockedStatus, Relation, Unit
 from pgconnstr import ConnectionString
 
-from constants import APP_SCOPE, DATABASE_PORT
+from constants import (
+    ALL_LEGACY_RELATIONS,
+    APP_SCOPE,
+    DATABASE_PORT,
+    ENDPOINT_SIMULTANEOUSLY_BLOCKING_MESSAGE,
+)
 from utils import new_password
 
 logger = logging.getLogger(__name__)
@@ -87,6 +92,20 @@ def _check_for_blocking_relations(self, relation_id: int) -> bool:
                         return True
         return False
 
+    def _check_exist_current_relation(self) -> bool:
+        for r in self.charm.client_relations:
+            if r in ALL_LEGACY_RELATIONS:
+                return True
+        return False
+
+    def _check_multiple_endpoints(self) -> bool:
+        """Checks if there are relations with other endpoints."""
+        is_exist = self._check_exist_current_relation()
+        for relation in self.charm.client_relations:
+            if relation.name not in ALL_LEGACY_RELATIONS and is_exist:
+                return True
+        return False
+
     def _on_relation_changed(self, event: RelationChangedEvent) -> None:
         """Handle the legacy db/db-admin relation changed event.
 
@@ -96,6 +115,10 @@ def _on_relation_changed(self, event: RelationChangedEvent) -> None:
         if not self.charm.unit.is_leader():
             return
 
+        if self._check_multiple_endpoints():
+            self.charm.unit.status = BlockedStatus(ENDPOINT_SIMULTANEOUSLY_BLOCKING_MESSAGE)
+            return
+
         if (
             "cluster_initialised" not in self.charm._peers.data[self.charm.app]
             or not self.charm._patroni.member_started
@@ -277,6 +300,16 @@ def _update_unit_status(self, relation: Relation) -> None:
         ]:
             if not self._check_for_blocking_relations(relation.id):
                 self.charm.unit.status = ActiveStatus()
+        self._update_unit_status_on_blocking_endpoint_simultaneously()
+
+    def _update_unit_status_on_blocking_endpoint_simultaneously(self):
+        """Clean up Blocked status if this is due related of multiple endpoints."""
+        if (
+            self.charm.is_blocked
+            and self.charm.unit.status.message == ENDPOINT_SIMULTANEOUSLY_BLOCKING_MESSAGE
+        ):
+            if not self._check_multiple_endpoints():
+                self.charm.unit.status = ActiveStatus()
 
     def update_endpoints(self, relation: Relation = None) -> None:
         """Set the read/write and read-only endpoints."""

src/relations/postgresql_provider.py

@@ -17,11 +17,16 @@
     PostgreSQLGetPostgreSQLVersionError,
     PostgreSQLListUsersError,
 )
-from ops.charm import CharmBase, RelationBrokenEvent
+from ops.charm import CharmBase, RelationBrokenEvent, RelationChangedEvent
 from ops.framework import Object
 from ops.model import ActiveStatus, BlockedStatus, Relation
 
-from constants import ALL_CLIENT_RELATIONS, APP_SCOPE, DATABASE_PORT
+from constants import (
+    ALL_CLIENT_RELATIONS,
+    APP_SCOPE,
+    DATABASE_PORT,
+    ENDPOINT_SIMULTANEOUSLY_BLOCKING_MESSAGE,
+)
 from utils import new_password
 
 logger = logging.getLogger(__name__)
@@ -48,7 +53,10 @@ def __init__(self, charm: CharmBase, relation_name: str = "database") -> None:
         self.framework.observe(
             charm.on[self.relation_name].relation_broken, self._on_relation_broken
         )
-
+        self.framework.observe(
+            charm.on[self.relation_name].relation_changed,
+            self._on_relation_changed_event,
+        )
         self.charm = charm
 
         # Charm events defined in the database provides charm library.
@@ -190,6 +198,13 @@ def update_endpoints(self, event: DatabaseRequestedEvent = None) -> None:
                 read_only_endpoints,
             )
 
+    def _check_multiple_endpoints(self) -> bool:
+        """Checks if there are relations with other endpoints."""
+        relation_names = {relation.name for relation in self.charm.client_relations}
+        if "database" in relation_names and len(relation_names) > 1:
+            return True
+        return False
+
     def _update_unit_status(self, relation: Relation) -> None:
         """# Clean up Blocked status if it's due to extensions request."""
         if (
@@ -199,6 +214,27 @@ def _update_unit_status(self, relation: Relation) -> None:
             if not self.check_for_invalid_extra_user_roles(relation.id):
                 self.charm.unit.status = ActiveStatus()
 
+        self._update_unit_status_on_blocking_endpoint_simultaneously()
+
+    def _on_relation_changed_event(self, event: RelationChangedEvent) -> None:
+        """Event emitted when the relation has changed."""
+        # Leader only
+        if not self.charm.unit.is_leader():
+            return
+
+        if self._check_multiple_endpoints():
+            self.charm.unit.status = BlockedStatus(ENDPOINT_SIMULTANEOUSLY_BLOCKING_MESSAGE)
+            return
+
+    def _update_unit_status_on_blocking_endpoint_simultaneously(self):
+        """Clean up Blocked status if this is due related of multiple endpoints."""
+        if (
+            self.charm.is_blocked
+            and self.charm.unit.status.message == ENDPOINT_SIMULTANEOUSLY_BLOCKING_MESSAGE
+        ):
+            if not self._check_multiple_endpoints():
+                self.charm.unit.status = ActiveStatus()
+
     def check_for_invalid_extra_user_roles(self, relation_id: int) -> bool:
         """Checks if there are relations with invalid extra user roles.
 

tests/integration/relations/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.

tests/integration/relations/helpers.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+from typing import Optional
+
+import yaml
+from pytest_operator.plugin import OpsTest
+
+
+async def get_legacy_db_connection_str(
+    ops_test: OpsTest,
+    application_name: str,
+    relation_name: str,
+    read_only_endpoint: bool = False,
+    remote_unit_name: str = None,
+) -> Optional[str]:
+    """Returns a PostgreSQL connection string.
+
+    Args:
+        ops_test: The ops test framework instance
+        application_name: The name of the application
+        relation_name: name of the relation to get connection data from
+        read_only_endpoint: whether to choose the read-only endpoint
+            instead of the read/write endpoint
+        remote_unit_name: Optional remote unit name used to retrieve
+            unit data instead of application data
+
+    Returns:
+        a PostgreSQL connection string
+    """
+    unit_name = f"{application_name}/0"
+    raw_data = (await ops_test.juju("show-unit", unit_name))[1]
+    if not raw_data:
+        raise ValueError(f"no unit info could be grabbed for {unit_name}")
+    data = yaml.safe_load(raw_data)
+    # Filter the data based on the relation name.
+    relation_data = [
+        v for v in data[unit_name]["relation-info"] if v["related-endpoint"] == relation_name
+    ]
+    if len(relation_data) == 0:
+        raise ValueError(
+            f"no relation data could be grabbed on relation with endpoint {relation_name}"
+        )
+    if remote_unit_name:
+        data = relation_data[0]["related-units"][remote_unit_name]["data"]
+    else:
+        data = relation_data[0]["application-data"]
+    if read_only_endpoint:
+        if data.get("standbys") is None:
+            return None
+        return data.get("standbys").split(",")[0]
+    else:
+        return data.get("master")

tests/integration/relations/test_relations.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+import asyncio
+import logging
+
+import psycopg2
+import pytest
+from pytest_operator.plugin import OpsTest
+from tenacity import Retrying, stop_after_delay, wait_fixed
+
+from ..helpers import CHARM_SERIES, METADATA
+from ..new_relations.test_new_relations import APPLICATION_APP_NAME, build_connection_string
+from ..relations.helpers import get_legacy_db_connection_str
+
+logger = logging.getLogger(__name__)
+
+APP_NAME = METADATA["name"]
+# MAILMAN3_CORE_APP_NAME = "mailman3-core"
+DB_RELATION = "db"
+DATABASE_RELATION = "database"
+FIRST_DATABASE_RELATION = "first-database"
+DATABASE_APP_NAME = "database-app"
+DB_APP_NAME = "db-app"
+APP_NAMES = [APP_NAME, DATABASE_APP_NAME, DB_APP_NAME]
+
+
+@pytest.mark.group(1)
+@pytest.mark.abort_on_fail
+async def test_deploy_charms(ops_test: OpsTest, charm):
+    """Deploy both charms (application and database) to use in the tests."""
+    # Deploy both charms (multiple units for each application to test that later they correctly
+    # set data in the relation application databag using only the leader unit).
+    async with ops_test.fast_forward():
+        await asyncio.gather(
+            ops_test.model.deploy(
+                APPLICATION_APP_NAME,
+                application_name=DATABASE_APP_NAME,
+                num_units=1,
+                series=CHARM_SERIES,
+                channel="edge",
+            ),
+            ops_test.model.deploy(
+                charm,
+                application_name=APP_NAME,
+                num_units=1,
+                series=CHARM_SERIES,
+                config={
+                    "profile": "testing",
+                    "plugin_unaccent_enable": "True",
+                    "plugin_pg_trgm_enable": "True",
+                },
+            ),
+            ops_test.model.deploy(
+                APPLICATION_APP_NAME,
+                application_name=DB_APP_NAME,
+                num_units=1,
+                series=CHARM_SERIES,
+                channel="edge",
+            ),
+        )
+
+        await ops_test.model.wait_for_idle(apps=APP_NAMES, status="active", timeout=3000)
+
+
+@pytest.mark.group(1)
+async def test_legacy_endpoint_with_multiple_related_endpoints(ops_test: OpsTest):
+    await ops_test.model.relate(f"{DB_APP_NAME}:{DB_RELATION}", f"{APP_NAME}:{DB_RELATION}")
+    await ops_test.model.relate(APP_NAME, f"{DATABASE_APP_NAME}:{FIRST_DATABASE_RELATION}")
+
+    app = ops_test.model.applications[APP_NAME]
+    await ops_test.model.block_until(
+        lambda: "blocked" in {unit.workload_status for unit in app.units},
+        timeout=1500,
+    )
+
+    logger.info(" remove relation with  modern endpoints")
+    await ops_test.model.applications[APP_NAME].remove_relation(
+        f"{APP_NAME}:{DATABASE_RELATION}", f"{DATABASE_APP_NAME}:{FIRST_DATABASE_RELATION}"
+    )
+    async with ops_test.fast_forward():
+        await ops_test.model.wait_for_idle(
+            status="active",
+            timeout=1500,
+            raise_on_error=False,
+        )
+
+    legacy_interface_connect = await get_legacy_db_connection_str(
+        ops_test, DB_APP_NAME, DB_RELATION, remote_unit_name=f"{APP_NAME}/0"
+    )
+    logger.info(f" check connect to = {legacy_interface_connect}")
+    for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(10)):
+        with attempt:
+            with psycopg2.connect(legacy_interface_connect) as connection:
+                assert connection.status == psycopg2.extensions.STATUS_READY
+
+    logger.info(f" remove relation {DB_APP_NAME}:{DB_RELATION}")
+    async with ops_test.fast_forward():
+        await ops_test.model.applications[APP_NAME].remove_relation(
+            f"{APP_NAME}:{DB_RELATION}", f"{DB_APP_NAME}:{DB_RELATION}"
+        )
+        await ops_test.model.wait_for_idle(apps=[APP_NAME], status="active", timeout=1000)
+        for attempt in Retrying(stop=stop_after_delay(60 * 5), wait=wait_fixed(10)):
+            with attempt:
+                with pytest.raises(psycopg2.OperationalError):
+                    psycopg2.connect(legacy_interface_connect)
+
+
+@pytest.mark.group(1)
+async def test_modern_endpoint_with_multiple_related_endpoints(ops_test: OpsTest):
+    await ops_test.model.relate(f"{DB_APP_NAME}:{DB_RELATION}", f"{APP_NAME}:{DB_RELATION}")
+    await ops_test.model.relate(APP_NAME, f"{DATABASE_APP_NAME}:{FIRST_DATABASE_RELATION}")
+
+    app = ops_test.model.applications[APP_NAME]
+    await ops_test.model.block_until(
+        lambda: "blocked" in {unit.workload_status for unit in app.units},
+        timeout=1500,
+    )
+
+    logger.info(" remove relation with legacy endpoints")
+    await ops_test.model.applications[APP_NAME].remove_relation(
+        f"{DB_APP_NAME}:{DB_RELATION}", f"{APP_NAME}:{DB_RELATION}"
+    )
+    async with ops_test.fast_forward():
+        await ops_test.model.wait_for_idle(status="active", timeout=3000, raise_on_error=False)
+
+    modern_interface_connect = await build_connection_string(
+        ops_test, DATABASE_APP_NAME, FIRST_DATABASE_RELATION
+    )
+    logger.info(f"check connect to = {modern_interface_connect}")
+    for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(10)):
+        with attempt:
+            with psycopg2.connect(modern_interface_connect) as connection:
+                assert connection.status == psycopg2.extensions.STATUS_READY
+
+    logger.info(f"remove relation {DATABASE_APP_NAME}:{FIRST_DATABASE_RELATION}")
+    async with ops_test.fast_forward():
+        await ops_test.model.applications[APP_NAME].remove_relation(
+            f"{APP_NAME}:{DATABASE_RELATION}", f"{DATABASE_APP_NAME}:{FIRST_DATABASE_RELATION}"
+        )
+        await ops_test.model.wait_for_idle(apps=[APP_NAME], status="active", timeout=1000)
+        for attempt in Retrying(stop=stop_after_delay(60 * 5), wait=wait_fixed(10)):
+            with attempt:
+                with pytest.raises(psycopg2.OperationalError):
+                    psycopg2.connect(modern_interface_connect)

tox.ini

@@ -26,6 +26,7 @@ allowlist_externals =
     charmcraft
     charmcraftcache
     mv
+    psycopg2-binary
 commands_pre =
     poetry export --only main,charm-libs --output requirements.txt
 commands =