[DPE-5387] Grants priviledges to non-public schemas (#742)

* grant access to non-public schemas * add unit tests
canonical · Oct 18, 2024 · 9ea543b · 9ea543b
1 parent 990076b
commit 9ea543b
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 10 deletions.
diff --git a/lib/charms/postgresql_k8s/v0/postgresql.py b/lib/charms/postgresql_k8s/v0/postgresql.py
@@ -36,7 +36,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 36
+LIBPATCH = 37
 
 INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE = "invalid role(s) for extra user roles"
 
@@ -393,24 +393,32 @@ def _generate_database_privileges_statements(
 SET lomowner = (SELECT oid FROM pg_roles WHERE rolname = '{}')
 WHERE lomowner = (SELECT oid FROM pg_roles WHERE rolname = '{}');""".format(user, self.user)
             )
+            for schema in schemas:
+                statements.append(
+                    sql.SQL("ALTER SCHEMA {} OWNER TO {};").format(
+                        sql.Identifier(schema), sql.Identifier(user)
+                    )
+                )
         else:
             for schema in schemas:
                 schema = sql.Identifier(schema)
-                statements.append(
+                statements.extend([
                     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)
-                    )
-                )
+                    ),
+                    sql.SQL("GRANT USAGE ON SCHEMA {} TO {};").format(
+                        schema, sql.Identifier(user)
+                    ),
+                    sql.SQL("GRANT CREATE ON SCHEMA {} TO {};").format(
+                        schema, sql.Identifier(user)
+                    ),
+                ])
         return statements
 
     def get_last_archived_wal(self) -> str:

diff --git a/tests/unit/test_postgresql.py b/tests/unit/test_postgresql.py
@@ -185,6 +185,20 @@ def test_generate_database_privileges_statements(harness):
             ),
         ]),
         "UPDATE pg_catalog.pg_largeobject_metadata\nSET lomowner = (SELECT oid FROM pg_roles WHERE rolname = 'test_user')\nWHERE lomowner = (SELECT oid FROM pg_roles WHERE rolname = 'operator');",
+        Composed([
+            SQL("ALTER SCHEMA "),
+            Identifier("test_schema_1"),
+            SQL(" OWNER TO "),
+            Identifier("test_user"),
+            SQL(";"),
+        ]),
+        Composed([
+            SQL("ALTER SCHEMA "),
+            Identifier("test_schema_2"),
+            SQL(" OWNER TO "),
+            Identifier("test_user"),
+            SQL(";"),
+        ]),
     ]
     # Test with multiple established relations.
     assert harness.charm.postgresql._generate_database_privileges_statements(
@@ -211,6 +225,20 @@ def test_generate_database_privileges_statements(harness):
             Identifier("test_user"),
             SQL(";"),
         ]),
+        Composed([
+            SQL("GRANT USAGE ON SCHEMA "),
+            Identifier("test_schema_1"),
+            SQL(" TO "),
+            Identifier("test_user"),
+            SQL(";"),
+        ]),
+        Composed([
+            SQL("GRANT CREATE ON SCHEMA "),
+            Identifier("test_schema_1"),
+            SQL(" TO "),
+            Identifier("test_user"),
+            SQL(";"),
+        ]),
         Composed([
             SQL("GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA "),
             Identifier("test_schema_2"),
@@ -232,6 +260,20 @@ def test_generate_database_privileges_statements(harness):
             Identifier("test_user"),
             SQL(";"),
         ]),
+        Composed([
+            SQL("GRANT USAGE ON SCHEMA "),
+            Identifier("test_schema_2"),
+            SQL(" TO "),
+            Identifier("test_user"),
+            SQL(";"),
+        ]),
+        Composed([
+            SQL("GRANT CREATE ON SCHEMA "),
+            Identifier("test_schema_2"),
+            SQL(" TO "),
+            Identifier("test_user"),
+            SQL(";"),
+        ]),
     ]