Fixed databases access (#198)

Author: marceloneppel

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -19,7 +19,7 @@
 Any charm using this library should import the `psycopg2` or `psycopg2-binary` dependency.
 """
 import logging
-from typing import Set
+from typing import List, Set
 
 import psycopg2
 from psycopg2 import sql
@@ -32,7 +32,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 9
+LIBPATCH = 10
 
 
 logger = logging.getLogger(__name__)
@@ -46,6 +46,10 @@ class PostgreSQLCreateUserError(Exception):
     """Exception raised when creating a user fails."""
 
 
+class PostgreSQLDatabasesSetupError(Exception):
+    """Exception raised when the databases setup fails."""
+
+
 class PostgreSQLDeleteUserError(Exception):
     """Exception raised when deleting a user fails."""
 
@@ -76,12 +80,14 @@ def __init__(
         user: str,
         password: str,
         database: str,
+        system_users: List[str] = [],
     ):
         self.primary_host = primary_host
         self.current_host = current_host
         self.user = user
         self.password = password
         self.database = database
+        self.system_users = system_users
 
     def _connect_to_database(
         self, database: str = None, connect_to_current_host: bool = False
@@ -119,10 +125,16 @@ def create_database(self, database: str, user: str) -> None:
             if cursor.fetchone() is None:
                 cursor.execute(sql.SQL("CREATE DATABASE {};").format(sql.Identifier(database)))
             cursor.execute(
-                sql.SQL("GRANT ALL PRIVILEGES ON DATABASE {} TO {};").format(
-                    sql.Identifier(database), sql.Identifier(user)
+                sql.SQL("REVOKE ALL PRIVILEGES ON DATABASE {} FROM PUBLIC;").format(
+                    sql.Identifier(database)
                 )
             )
+            for user_to_grant_access in [user] + self.system_users:
+                cursor.execute(
+                    sql.SQL("GRANT ALL PRIVILEGES ON DATABASE {} TO {};").format(
+                        sql.Identifier(database), sql.Identifier(user_to_grant_access)
+                    )
+                )
             with self._connect_to_database(database=database) as conn:
                 with conn.cursor() as curs:
                     statements = []
@@ -331,6 +343,26 @@ def list_users(self) -> Set[str]:
             logger.error(f"Failed to list PostgreSQL database users: {e}")
             raise PostgreSQLListUsersError()
 
+    def set_up_database(self) -> None:
+        """Set up postgres database with the right permissions."""
+        connection = None
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                # Allow access to the postgres database only to the system users.
+                cursor.execute("REVOKE ALL PRIVILEGES ON DATABASE postgres FROM PUBLIC;")
+                for user in self.system_users:
+                    cursor.execute(
+                        sql.SQL("GRANT ALL PRIVILEGES ON DATABASE postgres TO {};").format(
+                            sql.Identifier(user)
+                        )
+                    )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to set up databases: {e}")
+            raise PostgreSQLDatabasesSetupError()
+        finally:
+            if connection is not None:
+                connection.close()
+
     def update_user_password(self, username: str, password: str) -> None:
         """Update a user password.
 

src/charm.py

@@ -193,6 +193,7 @@ def postgresql(self) -> PostgreSQL:
             user=USER,
             password=self.get_secret("app", f"{USER}-password"),
             database="postgres",
+            system_users=SYSTEM_USERS,
         )
 
     @property
@@ -569,6 +570,8 @@ def _initialize_cluster(self, event: WorkloadEvent) -> bool:
                 extra_user_roles="pg_monitor",
             )
 
+        self.postgresql.set_up_database()
+
         # Mark the cluster as initialised.
         self._peers.data[self.app]["cluster_initialised"] = "True"
 

tests/integration/new_relations/helpers.py

@@ -18,6 +18,7 @@ async def build_connection_string(
     relation_id: str = None,
     relation_alias: str = None,
     read_only_endpoint: bool = False,
+    database: str = None,
 ) -> str:
     """Build a PostgreSQL connection string.
 
@@ -30,12 +31,14 @@ async def build_connection_string(
             to get connection data from
         read_only_endpoint: whether to choose the read-only endpoint
             instead of the read/write endpoint
+        database: optional database to be used in the connection string
 
     Returns:
         a PostgreSQL connection string
     """
     # Get the connection data exposed to the application through the relation.
-    database = f'{application_name.replace("-", "_")}_{relation_name.replace("-", "_")}'
+    if database is None:
+        database = f'{application_name.replace("-", "_")}_{relation_name.replace("-", "_")}'
     username = await get_application_relation_data(
         ops_test, application_name, relation_name, "username", relation_id, relation_alias
     )

tests/integration/new_relations/test_new_relations.py

@@ -180,6 +180,25 @@ async def test_two_applications_doesnt_share_the_same_relation_data(
 
     assert application_connection_string != another_application_connection_string
 
+    # Check that the user cannot access other databases.
+    for application, other_application_database in [
+        (APPLICATION_APP_NAME, "another_application_first_database"),
+        (another_application_app_name, "application_first_database"),
+    ]:
+        connection_string = await build_connection_string(
+            ops_test, application, FIRST_DATABASE_RELATION_NAME, database="postgres"
+        )
+        with pytest.raises(psycopg2.Error):
+            psycopg2.connect(connection_string)
+        connection_string = await build_connection_string(
+            ops_test,
+            application,
+            FIRST_DATABASE_RELATION_NAME,
+            database=other_application_database,
+        )
+        with pytest.raises(psycopg2.Error):
+            psycopg2.connect(connection_string)
+
 
 async def test_an_application_can_connect_to_multiple_database_clusters(
     ops_test: OpsTest, database_charm