Skip to content

Commit 3ceaeeb

Browse files
marceloneppelmarceloneppel
authored
Handle tables ownership (#298)
Signed-off-by: Marcelo Henrique Neppel <[email protected]>
1 parent c166e17 commit 3ceaeeb

File tree

7 files changed

+111
-26
lines changed

7 files changed

+111
-26
lines changed

‎lib/charms/postgresql_k8s/v0/postgresql.py

Lines changed: 69 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
from typing import Dict, List, Optional, Set, Tuple
2323

2424
import psycopg2
25+
from ops.model import Relation
2526
from psycopg2 import sql
27+
from psycopg2.sql import Composed
2628

2729
# The unique Charmhub library identifier, never change it
2830
LIBID = "24ee217a54e840a598ff21a079c3e678"
@@ -32,7 +34,7 @@
3234

3335
# Increment this PATCH version before using `charmcraft publish-lib` or reset
3436
# to 0 if you are raising the major API version
35-
LIBPATCH = 20
37+
LIBPATCH = 21
3638

3739
INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE = "invalid role(s) for extra user roles"
3840

@@ -117,13 +119,20 @@ def _connect_to_database(
117119
connection.autocommit = True
118120
return connection
119121

120-
def create_database(self, database: str, user: str, plugins: List[str] = []) -> None:
122+
def create_database(
123+
self,
124+
database: str,
125+
user: str,
126+
plugins: List[str] = [],
127+
client_relations: List[Relation] = [],
128+
) -> None:
121129
"""Creates a new database and grant privileges to a user on it.
122130
123131
Args:
124132
database: database to be created.
125133
user: user that will have access to the database.
126134
plugins: extensions to enable in the new database.
135+
client_relations: current established client relations.
127136
"""
128137
try:
129138
connection = self._connect_to_database()
@@ -142,29 +151,20 @@ def create_database(self, database: str, user: str, plugins: List[str] = []) ->
142151
sql.Identifier(database), sql.Identifier(user_to_grant_access)
143152
)
144153
)
154+
relations_accessing_this_database = 0
155+
for relation in client_relations:
156+
for data in relation.data.values():
157+
if data.get("database") == database:
158+
relations_accessing_this_database += 1
145159
with self._connect_to_database(database=database) as conn:
146160
with conn.cursor() as curs:
147-
statements = []
148161
curs.execute(
149162
"SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT LIKE 'pg_%' and schema_name <> 'information_schema';"
150163
)
151-
for row in curs:
152-
schema = sql.Identifier(row[0])
153-
statements.append(
154-
sql.SQL(
155-
"GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA {} TO {};"
156-
).format(schema, sql.Identifier(user))
157-
)
158-
statements.append(
159-
sql.SQL(
160-
"GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA {} TO {};"
161-
).format(schema, sql.Identifier(user))
162-
)
163-
statements.append(
164-
sql.SQL(
165-
"GRANT ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA {} TO {};"
166-
).format(schema, sql.Identifier(user))
167-
)
164+
schemas = [row[0] for row in curs.fetchall()]
165+
statements = self._generate_database_privileges_statements(
166+
relations_accessing_this_database, schemas, user
167+
)
168168
for statement in statements:
169169
curs.execute(statement)
170170
except psycopg2.Error as e:
@@ -308,6 +308,55 @@ def enable_disable_extensions(self, extensions: Dict[str, bool], database: str =
308308
if connection is not None:
309309
connection.close()
310310

311+
def _generate_database_privileges_statements(
312+
self, relations_accessing_this_database: int, schemas: List[str], user: str
313+
) -> List[Composed]:
314+
"""Generates a list of databases privileges statements."""
315+
statements = []
316+
if relations_accessing_this_database == 1:
317+
statements.append(
318+
sql.SQL(
319+
"""DO $$
320+
DECLARE r RECORD;
321+
BEGIN
322+
FOR r IN (SELECT statement FROM (SELECT 1 AS index,'ALTER TABLE '|| schemaname || '."' || tablename ||'" OWNER TO {};' AS statement
323+
FROM pg_tables WHERE NOT schemaname IN ('pg_catalog', 'information_schema')
324+
UNION SELECT 2 AS index,'ALTER SEQUENCE '|| sequence_schema || '."' || sequence_name ||'" OWNER TO {};' AS statement
325+
FROM information_schema.sequences WHERE NOT sequence_schema IN ('pg_catalog', 'information_schema')
326+
UNION SELECT 3 AS index,'ALTER FUNCTION '|| nsp.nspname || '."' || p.proname ||'"('||pg_get_function_identity_arguments(p.oid)||') OWNER TO {};' AS statement
327+
FROM pg_proc p JOIN pg_namespace nsp ON p.pronamespace = nsp.oid WHERE NOT nsp.nspname IN ('pg_catalog', 'information_schema')
328+
UNION SELECT 4 AS index,'ALTER VIEW '|| schemaname || '."' || viewname ||'" OWNER TO {};' AS statement
329+
FROM pg_catalog.pg_views WHERE NOT schemaname IN ('pg_catalog', 'information_schema')) AS statements ORDER BY index) LOOP
330+
EXECUTE format(r.statement);
331+
END LOOP;
332+
END; $$;"""
333+
).format(
334+
sql.Identifier(user),
335+
sql.Identifier(user),
336+
sql.Identifier(user),
337+
sql.Identifier(user),
338+
)
339+
)
340+
else:
341+
for schema in schemas:
342+
schema = sql.Identifier(schema)
343+
statements.append(
344+
sql.SQL("GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA {} TO {};").format(
345+
schema, sql.Identifier(user)
346+
)
347+
)
348+
statements.append(
349+
sql.SQL("GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA {} TO {};").format(
350+
schema, sql.Identifier(user)
351+
)
352+
)
353+
statements.append(
354+
sql.SQL("GRANT ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA {} TO {};").format(
355+
schema, sql.Identifier(user)
356+
)
357+
)
358+
return statements
359+
311360
def get_postgresql_text_search_configs(self) -> Set[str]:
312361
"""Returns the PostgreSQL available text search configs.
313362

‎src/charm.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1479,6 +1479,15 @@ def get_available_memory(self) -> int:
14791479

14801480
return 0
14811481

1482+
@property
1483+
def client_relations(self) -> List[Relation]:
1484+
"""Return the list of established client relations."""
1485+
relations = []
1486+
for relation_name in ["database", "db", "db-admin"]:
1487+
for relation in self.model.relations.get(relation_name, []):
1488+
relations.append(relation)
1489+
return relations
1490+
14821491

14831492
if __name__ == "__main__":
14841493
main(PostgresqlOperatorCharm)

‎src/relations/db.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,9 @@ def set_up_relation(self, relation: Relation) -> bool:
174174
if self.charm.config[plugin]
175175
]
176176

177-
self.charm.postgresql.create_database(database, user, plugins=plugins)
177+
self.charm.postgresql.create_database(
178+
database, user, plugins=plugins, client_relations=self.charm.client_relations
179+
)
178180

179181
except (PostgreSQLCreateDatabaseError, PostgreSQLCreateUserError) as e:
180182
logger.exception(e)

‎src/relations/postgresql_provider.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,9 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
9090
if self.charm.config[plugin]
9191
]
9292

93-
self.charm.postgresql.create_database(database, user, plugins=plugins)
93+
self.charm.postgresql.create_database(
94+
database, user, plugins=plugins, client_relations=self.charm.client_relations
95+
)
9496

9597
# Share the credentials with the application.
9698
self.database_provides.set_credentials(event.relation.id, user, password)

‎tests/unit/test_charm.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1538,3 +1538,18 @@ def test_juju_run_exec_divergence(self, _juju_version: Mock, _topology_observer:
15381538
harness = Harness(PostgresqlOperatorCharm)
15391539
harness.begin()
15401540
_topology_observer.assert_called_once_with(harness.charm, "/usr/bin/juju-exec")
1541+
1542+
def test_client_relations(self):
1543+
# Test when the charm has no relations.
1544+
self.assertEqual(self.charm.client_relations, [])
1545+
1546+
# Test when the charm has some relations.
1547+
self.harness.add_relation("database", "application")
1548+
self.harness.add_relation("db", "legacy-application")
1549+
self.harness.add_relation("db-admin", "legacy-admin-application")
1550+
database_relation = self.harness.model.get_relation("database")
1551+
db_relation = self.harness.model.get_relation("db")
1552+
db_admin_relation = self.harness.model.get_relation("db-admin")
1553+
self.assertEqual(
1554+
self.charm.client_relations, [database_relation, db_relation, db_admin_relation]
1555+
)

‎tests/unit/test_db.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -235,7 +235,9 @@ def test_set_up_relation(
235235
self.assertTrue(self.harness.charm.legacy_db_relation.set_up_relation(relation))
236236
user = f"relation-{self.rel_id}"
237237
postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
238-
postgresql_mock.create_database.assert_called_once_with(DATABASE, user, plugins=[])
238+
postgresql_mock.create_database.assert_called_once_with(
239+
DATABASE, user, plugins=[], client_relations=[relation]
240+
)
239241
_update_endpoints.assert_called_once()
240242
_update_unit_status.assert_called_once()
241243
self.assertNotIsInstance(self.harness.model.unit.status, BlockedStatus)
@@ -260,7 +262,9 @@ def test_set_up_relation(
260262
)
261263
self.assertTrue(self.harness.charm.legacy_db_relation.set_up_relation(relation))
262264
postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
263-
postgresql_mock.create_database.assert_called_once_with(DATABASE, user, plugins=[])
265+
postgresql_mock.create_database.assert_called_once_with(
266+
DATABASE, user, plugins=[], client_relations=[relation]
267+
)
264268
_update_endpoints.assert_called_once()
265269
_update_unit_status.assert_called_once()
266270
self.assertNotIsInstance(self.harness.model.unit.status, BlockedStatus)
@@ -280,7 +284,7 @@ def test_set_up_relation(
280284
self.assertTrue(self.harness.charm.legacy_db_relation.set_up_relation(relation))
281285
postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
282286
postgresql_mock.create_database.assert_called_once_with(
283-
"application", user, plugins=[]
287+
"application", user, plugins=[], client_relations=[relation]
284288
)
285289
_update_endpoints.assert_called_once()
286290
_update_unit_status.assert_called_once()

‎tests/unit/test_postgresql_provider.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,11 @@ def test_on_database_requested(
124124
postgresql_mock.create_user.assert_called_once_with(
125125
user, "test-password", extra_user_roles=EXTRA_USER_ROLES
126126
)
127-
postgresql_mock.create_database.assert_called_once_with(DATABASE, user, plugins=[])
127+
database_relation = self.harness.model.get_relation(RELATION_NAME)
128+
client_relations = [database_relation]
129+
postgresql_mock.create_database.assert_called_once_with(
130+
DATABASE, user, plugins=[], client_relations=client_relations
131+
)
128132
postgresql_mock.get_postgresql_version.assert_called_once()
129133
_update_endpoints.assert_called_once()
130134

0 commit comments

Comments
 (0)