Handle tables ownership (#334)

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

Author: marceloneppel

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -22,7 +22,9 @@
 from typing import Dict, List, Optional, Set, Tuple
 
 import psycopg2
+from ops.model import Relation
 from psycopg2 import sql
+from psycopg2.sql import Composed
 
 # The unique Charmhub library identifier, never change it
 LIBID = "24ee217a54e840a598ff21a079c3e678"
@@ -32,7 +34,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 20
+LIBPATCH = 21
 
 INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE = "invalid role(s) for extra user roles"
 
@@ -117,13 +119,20 @@ def _connect_to_database(
         connection.autocommit = True
         return connection
 
-    def create_database(self, database: str, user: str, plugins: List[str] = []) -> None:
+    def create_database(
+        self,
+        database: str,
+        user: str,
+        plugins: List[str] = [],
+        client_relations: List[Relation] = [],
+    ) -> None:
         """Creates a new database and grant privileges to a user on it.
 
         Args:
             database: database to be created.
             user: user that will have access to the database.
             plugins: extensions to enable in the new database.
+            client_relations: current established client relations.
         """
         try:
             connection = self._connect_to_database()
@@ -142,29 +151,20 @@ def create_database(self, database: str, user: str, plugins: List[str] = []) ->
                         sql.Identifier(database), sql.Identifier(user_to_grant_access)
                     )
                 )
+            relations_accessing_this_database = 0
+            for relation in client_relations:
+                for data in relation.data.values():
+                    if data.get("database") == database:
+                        relations_accessing_this_database += 1
             with self._connect_to_database(database=database) as conn:
                 with conn.cursor() as curs:
-                    statements = []
                     curs.execute(
                         "SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT LIKE 'pg_%' and schema_name <> 'information_schema';"
                     )
-                    for row in curs:
-                        schema = sql.Identifier(row[0])
-                        statements.append(
-                            sql.SQL(
-                                "GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA {} TO {};"
-                            ).format(schema, sql.Identifier(user))
-                        )
-                        statements.append(
-                            sql.SQL(
-                                "GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA {} TO {};"
-                            ).format(schema, sql.Identifier(user))
-                        )
-                        statements.append(
-                            sql.SQL(
-                                "GRANT ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA {} TO {};"
-                            ).format(schema, sql.Identifier(user))
-                        )
+                    schemas = [row[0] for row in curs.fetchall()]
+                    statements = self._generate_database_privileges_statements(
+                        relations_accessing_this_database, schemas, user
+                    )
                     for statement in statements:
                         curs.execute(statement)
         except psycopg2.Error as e:
@@ -308,6 +308,55 @@ def enable_disable_extensions(self, extensions: Dict[str, bool], database: str =
             if connection is not None:
                 connection.close()
 
+    def _generate_database_privileges_statements(
+        self, relations_accessing_this_database: int, schemas: List[str], user: str
+    ) -> List[Composed]:
+        """Generates a list of databases privileges statements."""
+        statements = []
+        if relations_accessing_this_database == 1:
+            statements.append(
+                sql.SQL(
+                    """DO $$
+DECLARE r RECORD;
+BEGIN
+  FOR r IN (SELECT statement FROM (SELECT 1 AS index,'ALTER TABLE '|| schemaname || '."' || tablename ||'" OWNER TO {};' AS statement
+FROM pg_tables WHERE NOT schemaname IN ('pg_catalog', 'information_schema')
+UNION SELECT 2 AS index,'ALTER SEQUENCE '|| sequence_schema || '."' || sequence_name ||'" OWNER TO {};' AS statement
+FROM information_schema.sequences WHERE NOT sequence_schema IN ('pg_catalog', 'information_schema')
+UNION SELECT 3 AS index,'ALTER FUNCTION '|| nsp.nspname || '."' || p.proname ||'"('||pg_get_function_identity_arguments(p.oid)||') OWNER TO {};' AS statement
+FROM pg_proc p JOIN pg_namespace nsp ON p.pronamespace = nsp.oid WHERE NOT nsp.nspname IN ('pg_catalog', 'information_schema')
+UNION SELECT 4 AS index,'ALTER VIEW '|| schemaname || '."' || viewname ||'" OWNER TO {};' AS statement
+FROM pg_catalog.pg_views WHERE NOT schemaname IN ('pg_catalog', 'information_schema')) AS statements ORDER BY index) LOOP
+      EXECUTE format(r.statement);
+  END LOOP;
+END; $$;"""
+                ).format(
+                    sql.Identifier(user),
+                    sql.Identifier(user),
+                    sql.Identifier(user),
+                    sql.Identifier(user),
+                )
+            )
+        else:
+            for schema in schemas:
+                schema = sql.Identifier(schema)
+                statements.append(
+                    sql.SQL("GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA {} TO {};").format(
+                        schema, sql.Identifier(user)
+                    )
+                )
+                statements.append(
+                    sql.SQL("GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA {} TO {};").format(
+                        schema, sql.Identifier(user)
+                    )
+                )
+                statements.append(
+                    sql.SQL("GRANT ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA {} TO {};").format(
+                        schema, sql.Identifier(user)
+                    )
+                )
+        return statements
+
     def get_postgresql_text_search_configs(self) -> Set[str]:
         """Returns the PostgreSQL available text search configs.
 

src/charm.py

@@ -1512,6 +1512,15 @@ def get_available_resources(self) -> Tuple[int, int]:
 
         return cpu_cores, allocable_memory
 
+    @property
+    def client_relations(self) -> List[Relation]:
+        """Return the list of established client relations."""
+        relations = []
+        for relation_name in ["database", "db", "db-admin"]:
+            for relation in self.model.relations.get(relation_name, []):
+                relations.append(relation)
+        return relations
+
 
 if __name__ == "__main__":
     main(PostgresqlOperatorCharm, use_juju_for_storage=True)

src/relations/db.py

@@ -150,7 +150,9 @@ def set_up_relation(self, relation: Relation) -> bool:
                 if self.charm.config[plugin]
             ]
 
-            self.charm.postgresql.create_database(database, user, plugins=plugins)
+            self.charm.postgresql.create_database(
+                database, user, plugins=plugins, client_relations=self.charm.client_relations
+            )
 
             # Build the primary's connection string.
             primary = str(

src/relations/postgresql_provider.py

@@ -91,7 +91,9 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
                 if self.charm.config[plugin]
             ]
 
-            self.charm.postgresql.create_database(database, user, plugins=plugins)
+            self.charm.postgresql.create_database(
+                database, user, plugins=plugins, client_relations=self.charm.client_relations
+            )
 
             # Share the credentials with the application.
             self.database_provides.set_credentials(event.relation.id, user, password)

tests/integration/new_relations/test_new_relations.py

@@ -5,6 +5,7 @@
 import logging
 import secrets
 import string
+from asyncio import gather
 from pathlib import Path
 
 import psycopg2
@@ -30,6 +31,8 @@
 DATABASE_APP_NAME = "database"
 ANOTHER_DATABASE_APP_NAME = "another-database"
 DATA_INTEGRATOR_APP_NAME = "data-integrator"
+DISCOURSE_APP_NAME = "discourse-k8s"
+REDIS_APP_NAME = "redis-k8s"
 APP_NAMES = [APPLICATION_APP_NAME, DATABASE_APP_NAME, ANOTHER_DATABASE_APP_NAME]
 DATABASE_APP_METADATA = yaml.safe_load(Path("./metadata.yaml").read_text())
 FIRST_DATABASE_RELATION_NAME = "first-database"
@@ -85,7 +88,9 @@ async def test_database_relation_with_charm_libraries(ops_test: OpsTest, databas
         await ops_test.model.add_relation(
             f"{APPLICATION_APP_NAME}:{FIRST_DATABASE_RELATION_NAME}", DATABASE_APP_NAME
         )
-        await ops_test.model.wait_for_idle(apps=APP_NAMES, status="active", raise_on_blocked=True)
+        await ops_test.model.wait_for_idle(
+            apps=[DATABASE_APP_NAME], status="active", raise_on_blocked=True
+        )
 
         # Check that on juju 3 we have secrets and no username and password in the rel databag
         if hasattr(ops_test.model, "list_secrets"):
@@ -571,6 +576,81 @@ async def test_invalid_extra_user_roles(ops_test: OpsTest):
         )
 
 
+async def test_discourse(ops_test: OpsTest):
+    # Deploy Discourse and Redis.
+    await gather(
+        ops_test.model.deploy(DISCOURSE_APP_NAME, application_name=DISCOURSE_APP_NAME),
+        ops_test.model.deploy(
+            REDIS_APP_NAME, application_name=REDIS_APP_NAME, channel="latest/edge"
+        ),
+    )
+
+    async with ops_test.fast_forward():
+        # Enable the plugins/extensions required by Discourse.
+        logger.info("Enabling the plugins/extensions required by Discourse")
+        config = {"plugin_hstore_enable": "True", "plugin_pg_trgm_enable": "True"}
+        await ops_test.model.applications[DATABASE_APP_NAME].set_config(config)
+        await gather(
+            ops_test.model.wait_for_idle(apps=[DISCOURSE_APP_NAME], status="waiting"),
+            ops_test.model.wait_for_idle(
+                apps=[DATABASE_APP_NAME, REDIS_APP_NAME], status="active"
+            ),
+        )
+        # Add both relations to Discourse (PostgreSQL and Redis)
+        # and wait for it to be ready.
+        logger.info("Adding relations")
+        await gather(
+            ops_test.model.add_relation(DATABASE_APP_NAME, DISCOURSE_APP_NAME),
+            ops_test.model.add_relation(REDIS_APP_NAME, DISCOURSE_APP_NAME),
+        )
+        await gather(
+            ops_test.model.wait_for_idle(apps=[DISCOURSE_APP_NAME], timeout=2000),
+            ops_test.model.wait_for_idle(
+                apps=[DATABASE_APP_NAME, REDIS_APP_NAME], status="active"
+            ),
+        )
+        logger.info("Configuring Discourse")
+        config = {
+            "developer_emails": "[email protected]",
+            "external_hostname": "discourse-k8s",
+            "smtp_address": "test.local",
+            "smtp_domain": "test.local",
+            "s3_install_cors_rule": "false",
+        }
+        await ops_test.model.applications[DISCOURSE_APP_NAME].set_config(config)
+        await ops_test.model.wait_for_idle(apps=[DISCOURSE_APP_NAME], status="active")
+
+        # Deploy a new discourse application (https://github.com/canonical/data-platform-libs/issues/118
+        # prevents from re-relating the same Discourse application; Discourse uses the old secret and fails).
+        await ops_test.model.applications[DISCOURSE_APP_NAME].remove()
+        other_discourse_app_name = f"other-{DISCOURSE_APP_NAME}"
+        await ops_test.model.deploy(DISCOURSE_APP_NAME, application_name=other_discourse_app_name)
+
+        # Add both relations to Discourse (PostgreSQL and Redis)
+        # and wait for it to be ready.
+        logger.info("Adding relations")
+        await gather(
+            ops_test.model.add_relation(DATABASE_APP_NAME, other_discourse_app_name),
+            ops_test.model.add_relation(REDIS_APP_NAME, other_discourse_app_name),
+        )
+        await gather(
+            ops_test.model.wait_for_idle(apps=[other_discourse_app_name], timeout=2000),
+            ops_test.model.wait_for_idle(
+                apps=[DATABASE_APP_NAME, REDIS_APP_NAME], status="active"
+            ),
+        )
+        logger.info("Configuring Discourse")
+        config = {
+            "developer_emails": "[email protected]",
+            "external_hostname": "discourse-k8s",
+            "smtp_address": "test.local",
+            "smtp_domain": "test.local",
+            "s3_install_cors_rule": "false",
+        }
+        await ops_test.model.applications[other_discourse_app_name].set_config(config)
+        await ops_test.model.wait_for_idle(apps=[other_discourse_app_name], status="active")
+
+
 async def test_indico_datatabase(ops_test: OpsTest) -> None:
     """Tests deploying and relating to the Indico charm."""
     async with ops_test.fast_forward(fast_interval="30s"):

tests/integration/test_db.py

@@ -4,7 +4,6 @@
 import logging
 from asyncio import gather
 
-import pytest
 from pytest_operator.plugin import OpsTest
 
 from tests.integration.helpers import (
@@ -16,15 +15,12 @@
     check_database_users_existence,
     deploy_and_relate_application_with_postgresql,
     get_leader_unit,
-    get_primary,
     wait_for_relation_removed_between,
 )
 
 EXTENSIONS_BLOCKING_MESSAGE = "extensions requested through relation"
 FINOS_WALTZ_APP_NAME = "finos-waltz"
 ANOTHER_FINOS_WALTZ_APP_NAME = "another-finos-waltz"
-DISCOURSE_APP_NAME = "discourse-k8s"
-REDIS_APP_NAME = "redis-k8s"
 APPLICATION_UNITS = 1
 DATABASE_UNITS = 3
 
@@ -176,62 +172,3 @@ async def test_extensions_blocking(ops_test: OpsTest) -> None:
         raise_on_blocked=False,
         timeout=2000,
     )
-
-
-@pytest.mark.skip(reason="Should be ported and moved to the new relation tests")
-async def test_discourse(ops_test: OpsTest):
-    database_application_name = f"extensions-{DATABASE_APP_NAME}"
-    if database_application_name not in ops_test.model.applications:
-        await build_and_deploy(ops_test, 1, database_application_name)
-
-    # Deploy Discourse and Redis.
-    await gather(
-        ops_test.model.deploy(DISCOURSE_APP_NAME, application_name=DISCOURSE_APP_NAME),
-        ops_test.model.deploy(REDIS_APP_NAME, application_name=REDIS_APP_NAME),
-    )
-
-    # Add both relations to Discourse (PostgreSQL and Redis)
-    # and wait for it to be ready.
-    relation = await ops_test.model.add_relation(
-        f"{database_application_name}:db",
-        DISCOURSE_APP_NAME,
-    )
-    await ops_test.model.add_relation(
-        REDIS_APP_NAME,
-        DISCOURSE_APP_NAME,
-    )
-
-    # Discourse requests extensions through relation, so check that the PostgreSQL charm
-    # becomes blocked.
-    primary_name = await get_primary(ops_test, database_app_name=database_application_name)
-    async with ops_test.fast_forward():
-        await gather(
-            ops_test.model.block_until(
-                lambda: ops_test.model.units[primary_name].workload_status == "blocked", timeout=60
-            ),
-            ops_test.model.wait_for_idle(
-                apps=[REDIS_APP_NAME], status="active", timeout=2000
-            ),  # Redis sometimes takes a longer time to become active.
-        )
-        assert (
-            ops_test.model.units[primary_name].workload_status_message
-            == EXTENSIONS_BLOCKING_MESSAGE
-        )
-
-        # Enable the plugins/extensions required by Discourse.
-        config = {"plugin_hstore_enable": "True", "plugin_pg_trgm_enable": "True"}
-        await ops_test.model.applications[database_application_name].set_config(config)
-        await ops_test.model.wait_for_idle(
-            apps=[database_application_name, DISCOURSE_APP_NAME, REDIS_APP_NAME],
-            status="active",
-            timeout=2000,
-        )
-
-        # Check for the correct databases and users creation.
-        await check_database_creation(
-            ops_test, "discourse", database_app_name=database_application_name
-        )
-        discourse_users = [f"relation_id_{relation.id}"]
-        await check_database_users_existence(
-            ops_test, discourse_users, [], database_app_name=database_application_name
-        )

tests/unit/test_charm.py

@@ -716,6 +716,21 @@ def test_on_stop(self, _client):
             self.assertEqual(_client.return_value.apply.call_count, 2)
             self.assertIn("failed to patch k8s MagicMock", "".join(logs.output))
 
+    def test_client_relations(self):
+        # Test when the charm has no relations.
+        self.assertEqual(self.charm.client_relations, [])
+
+        # Test when the charm has some relations.
+        self.harness.add_relation("database", "application")
+        self.harness.add_relation("db", "legacy-application")
+        self.harness.add_relation("db-admin", "legacy-admin-application")
+        database_relation = self.harness.model.get_relation("database")
+        db_relation = self.harness.model.get_relation("db")
+        db_admin_relation = self.harness.model.get_relation("db-admin")
+        self.assertEqual(
+            self.charm.client_relations, [database_relation, db_relation, db_admin_relation]
+        )
+
     @parameterized.expand([("app"), ("unit")])
     @pytest.mark.usefixtures("only_with_juju_secrets")
     def test_set_secret_returning_secret_label(self, scope):