Skip to content

Commit 3ebd14c

Browse files
marceloneppelmarceloneppel
authored
[DPE-7584] Recreate temp tablespace on reboot (#11)
* Recreate temp tablespace on reboot Signed-off-by: Marcelo Henrique Neppel <[email protected]> * Fix unit test Signed-off-by: Marcelo Henrique Neppel <[email protected]> * Improve unit tests Signed-off-by: Marcelo Henrique Neppel <[email protected]> * Rename tablespace instead of dropping it Signed-off-by: Marcelo Henrique Neppel <[email protected]> * Change timestamp format Signed-off-by: Marcelo Henrique Neppel <[email protected]> * Log count and names of temp tablespaces Signed-off-by: Marcelo Henrique Neppel <[email protected]> --------- Signed-off-by: Marcelo Henrique Neppel <[email protected]>
1 parent dcea141 commit 3ebd14c

File tree

4 files changed

+124
-10
lines changed

4 files changed

+124
-10
lines changed

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
[project]
55
name = "postgresql-charms-single-kernel"
66
description = "Shared and reusable code for PostgreSQL-related charms"
7-
version = "16.0.0"
7+
version = "16.0.1"
88
readme = "README.md"
99
license = "Apache-2.0"
1010
authors = [

‎single_kernel_postgresql/utils/postgresql.py‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,16 @@
2121

2222
import logging
2323
import os
24+
import pwd
2425
from collections import OrderedDict
26+
from datetime import datetime, timezone
2527
from typing import Dict, List, Optional, Set, Tuple
2628

2729
import psycopg2
2830
from ops import ConfigData
2931
from psycopg2.sql import SQL, Identifier, Literal
3032

31-
from ..config.literals import BACKUP_USER, POSTGRESQL_STORAGE_PERMISSIONS, SYSTEM_USERS
33+
from ..config.literals import BACKUP_USER, POSTGRESQL_STORAGE_PERMISSIONS, SNAP_USER, SYSTEM_USERS
3234
from .filesystem import change_owner
3335

3436
# Groups to distinguish HBA access
@@ -1074,9 +1076,31 @@ def set_up_database(self, temp_location: Optional[str] = None) -> None:
10741076

10751077
if temp_location is not None:
10761078
# Fix permissions on the temporary tablespace location when a reboot happens and tmpfs is being used.
1077-
change_owner(temp_location)
1078-
os.chmod(temp_location, POSTGRESQL_STORAGE_PERMISSIONS)
1079+
temp_location_stats = os.stat(temp_location)
1080+
if (
1081+
pwd.getpwuid(temp_location_stats.st_uid).pw_name != SNAP_USER
1082+
or temp_location_stats.st_mode != POSTGRESQL_STORAGE_PERMISSIONS
1083+
):
1084+
change_owner(temp_location)
1085+
os.chmod(temp_location, POSTGRESQL_STORAGE_PERMISSIONS)
1086+
# Rename existing temp tablespace if it exists, instead of dropping it.
1087+
cursor.execute("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';")
1088+
if cursor.fetchone() is not None:
1089+
new_name = f"temp_{datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}"
1090+
cursor.execute(f"ALTER TABLESPACE temp RENAME TO {new_name};")
1091+
1092+
# List temp tablespaces with suffix for operator follow-up cleanup and log them.
1093+
cursor.execute(
1094+
"SELECT spcname FROM pg_tablespace WHERE spcname LIKE 'temp_%';"
1095+
)
1096+
temp_tbls = sorted([row[0] for row in cursor.fetchall()])
1097+
logger.info(
1098+
"There are %d temp tablespaces that should be checked and removed: %s",
1099+
len(temp_tbls),
1100+
", ".join(temp_tbls),
1101+
)
10791102

1103+
# Ensure a fresh temp tablespace exists at the expected location.
10801104
cursor.execute("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';")
10811105
if cursor.fetchone() is None:
10821106
cursor.execute(f"CREATE TABLESPACE temp LOCATION '{temp_location}';")

‎tests/unit/test_postgresql.py‎

Lines changed: 95 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,32 @@
11
# Copyright 2025 Canonical Ltd.
22
# See LICENSE file for licensing details.
3+
# ruff: noqa: I001
34
from unittest.mock import call, patch, sentinel
5+
from datetime import datetime, timezone, UTC
46

57
import psycopg2
68
import pytest
79
from ops.testing import Harness
8-
from psycopg2.sql import SQL, Composed, Identifier, Literal
10+
from psycopg2.sql import Composed, Identifier, Literal, SQL
11+
912
from single_kernel_postgresql.abstract_charm import AbstractPostgreSQLCharm
10-
from single_kernel_postgresql.config.literals import PEER, SYSTEM_USERS
13+
from single_kernel_postgresql.config.literals import (
14+
PEER,
15+
POSTGRESQL_STORAGE_PERMISSIONS,
16+
SNAP_USER,
17+
SYSTEM_USERS,
18+
)
1119
from single_kernel_postgresql.utils.postgresql import (
12-
ACCESS_GROUP_INTERNAL,
1320
ACCESS_GROUPS,
14-
ROLE_DATABASES_OWNER,
21+
ACCESS_GROUP_INTERNAL,
1522
PostgreSQL,
1623
PostgreSQLCreateDatabaseError,
1724
PostgreSQLCreateUserError,
1825
PostgreSQLDatabasesSetupError,
1926
PostgreSQLGetLastArchivedWALError,
2027
PostgreSQLUndefinedHostError,
2128
PostgreSQLUndefinedPasswordError,
29+
ROLE_DATABASES_OWNER,
2230
)
2331

2432

@@ -329,7 +337,14 @@ def test_set_up_database_with_temp_tablespace_and_missing_owner_role(harness):
329337
patch("single_kernel_postgresql.utils.postgresql.PostgreSQL.create_user") as _create_user,
330338
patch("single_kernel_postgresql.utils.postgresql.change_owner") as _change_owner,
331339
patch("single_kernel_postgresql.utils.postgresql.os.chmod") as _chmod,
340+
patch("single_kernel_postgresql.utils.postgresql.os.stat") as _stat,
341+
patch("single_kernel_postgresql.utils.postgresql.pwd.getpwuid") as _getpwuid,
332342
):
343+
# Simulate a temp location owned by wrong user/permissions to trigger fixup
344+
stat_result = type("stat_result", (), {"st_uid": 0, "st_mode": 0o755})
345+
_stat.return_value = stat_result
346+
_getpwuid.return_value.pw_name = "root"
347+
333348
# First connection (non-context) for temp tablespace
334349
execute_direct = _connect_to_database.return_value.cursor.return_value.execute
335350
fetchone_direct = _connect_to_database.return_value.cursor.return_value.fetchone
@@ -346,8 +361,9 @@ def test_set_up_database_with_temp_tablespace_and_missing_owner_role(harness):
346361
_change_owner.assert_called_once_with("/var/lib/postgresql/tmp")
347362
_chmod.assert_called_once_with("/var/lib/postgresql/tmp", 0o700)
348363

349-
# Validate temp tablespace operations
364+
# Validate temp tablespace operations: check existence and create/grant when missing
350365
execute_direct.assert_has_calls([
366+
call("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';"),
351367
call("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';"),
352368
call("CREATE TABLESPACE temp LOCATION '/var/lib/postgresql/tmp';"),
353369
call("GRANT CREATE ON TABLESPACE temp TO public;"),
@@ -371,6 +387,80 @@ def test_set_up_database_with_temp_tablespace_and_missing_owner_role(harness):
371387
execute_cm.assert_has_calls(expected, any_order=False)
372388

373389

390+
def test_set_up_database_owner_mismatch_triggers_rename_and_fix(harness):
391+
with (
392+
patch(
393+
"single_kernel_postgresql.utils.postgresql.PostgreSQL._connect_to_database"
394+
) as _connect_to_database,
395+
patch("single_kernel_postgresql.utils.postgresql.PostgreSQL.set_up_login_hook_function"),
396+
patch(
397+
"single_kernel_postgresql.utils.postgresql.PostgreSQL.set_up_predefined_catalog_roles_function"
398+
),
399+
patch("single_kernel_postgresql.utils.postgresql.change_owner") as _change_owner,
400+
patch("single_kernel_postgresql.utils.postgresql.os.chmod") as _chmod,
401+
patch("single_kernel_postgresql.utils.postgresql.os.stat") as _stat,
402+
patch("single_kernel_postgresql.utils.postgresql.pwd.getpwuid") as _getpwuid,
403+
patch("single_kernel_postgresql.utils.postgresql.datetime") as _dt,
404+
):
405+
# Owner differs, permissions are correct
406+
stat_result = type(
407+
"stat_result", (), {"st_uid": 0, "st_mode": POSTGRESQL_STORAGE_PERMISSIONS}
408+
)
409+
_stat.return_value = stat_result
410+
_getpwuid.return_value.pw_name = "root"
411+
412+
# Mock datetime.now(timezone.utc) to a fixed timestamp
413+
_dt.now.return_value = datetime(2025, 1, 1, 1, 2, 3, tzinfo=UTC)
414+
_dt.timezone = timezone # ensure timezone.utc is available in the patch target
415+
416+
execute_direct = _connect_to_database.return_value.cursor.return_value.execute
417+
fetchone_direct = _connect_to_database.return_value.cursor.return_value.fetchone
418+
fetchone_direct.side_effect = [True, None]
419+
420+
harness.charm.postgresql.set_up_database(temp_location="/var/lib/postgresql/tmp")
421+
422+
_change_owner.assert_called_once_with("/var/lib/postgresql/tmp")
423+
_chmod.assert_called_once_with("/var/lib/postgresql/tmp", POSTGRESQL_STORAGE_PERMISSIONS)
424+
execute_direct.assert_any_call("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';")
425+
execute_direct.assert_any_call("ALTER TABLESPACE temp RENAME TO temp_20250101010203;")
426+
427+
428+
def test_set_up_database_permissions_mismatch_triggers_rename_and_fix(harness):
429+
with (
430+
patch(
431+
"single_kernel_postgresql.utils.postgresql.PostgreSQL._connect_to_database"
432+
) as _connect_to_database,
433+
patch("single_kernel_postgresql.utils.postgresql.PostgreSQL.set_up_login_hook_function"),
434+
patch(
435+
"single_kernel_postgresql.utils.postgresql.PostgreSQL.set_up_predefined_catalog_roles_function"
436+
),
437+
patch("single_kernel_postgresql.utils.postgresql.change_owner") as _change_owner,
438+
patch("single_kernel_postgresql.utils.postgresql.os.chmod") as _chmod,
439+
patch("single_kernel_postgresql.utils.postgresql.os.stat") as _stat,
440+
patch("single_kernel_postgresql.utils.postgresql.pwd.getpwuid") as _getpwuid,
441+
patch("single_kernel_postgresql.utils.postgresql.datetime") as _dt,
442+
):
443+
# Owner matches SNAP_USER, permissions differ
444+
stat_result = type("stat_result", (), {"st_uid": 0, "st_mode": 0o755})
445+
_stat.return_value = stat_result
446+
_getpwuid.return_value.pw_name = SNAP_USER
447+
448+
# Mock datetime.now(timezone.utc) to a fixed timestamp
449+
_dt.now.return_value = datetime(2025, 1, 1, 1, 2, 3, tzinfo=UTC)
450+
_dt.timezone = timezone
451+
452+
execute_direct = _connect_to_database.return_value.cursor.return_value.execute
453+
fetchone_direct = _connect_to_database.return_value.cursor.return_value.fetchone
454+
fetchone_direct.side_effect = [True, None]
455+
456+
harness.charm.postgresql.set_up_database(temp_location="/var/lib/postgresql/tmp")
457+
458+
_change_owner.assert_called_once_with("/var/lib/postgresql/tmp")
459+
_chmod.assert_called_once_with("/var/lib/postgresql/tmp", POSTGRESQL_STORAGE_PERMISSIONS)
460+
execute_direct.assert_any_call("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';")
461+
execute_direct.assert_any_call("ALTER TABLESPACE temp RENAME TO temp_20250101010203;")
462+
463+
374464
def test_set_up_database_no_temp_and_existing_owner_role(harness):
375465
with (
376466
patch(

‎uv.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)