Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 10 additions & 12 deletions lib/charms/postgresql_k8s/v0/postgresql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = 19
LIBPATCH = 20

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

Expand Down Expand Up @@ -172,8 +172,7 @@ def create_database(self, database: str, user: str, plugins: List[str] = []) ->
raise PostgreSQLCreateDatabaseError()

# Enable preset extensions
for plugin in plugins:
self.enable_disable_extension(plugin, True, database)
self.enable_disable_extensions({plugin: True for plugin in plugins}, database)

def create_user(
self, user: str, password: str = None, admin: bool = False, extra_user_roles: str = None
Expand Down Expand Up @@ -270,22 +269,16 @@ def delete_user(self, user: str) -> None:
logger.error(f"Failed to delete user: {e}")
raise PostgreSQLDeleteUserError()

def enable_disable_extension(self, extension: str, enable: bool, database: str = None) -> None:
def enable_disable_extensions(self, extensions: Dict[str, bool], database: str = None) -> None:
"""Enables or disables a PostgreSQL extension.

Args:
extension: the name of the extensions.
enable: whether the extension should be enabled or disabled.
extensions: the name of the extensions.
database: optional database where to enable/disable the extension.

Raises:
PostgreSQLEnableDisableExtensionError if the operation fails.
"""
statement = (
f"CREATE EXTENSION IF NOT EXISTS {extension};"
if enable
else f"DROP EXTENSION IF EXISTS {extension};"
)
connection = None
try:
if database is not None:
Expand All @@ -301,7 +294,12 @@ def enable_disable_extension(self, extension: str, enable: bool, database: str =
with self._connect_to_database(
database=database
) as connection, connection.cursor() as cursor:
cursor.execute(statement)
for extension, enable in extensions.items():
cursor.execute(
f"CREATE EXTENSION IF NOT EXISTS {extension};"
if enable
else f"DROP EXTENSION IF EXISTS {extension};"
)
except psycopg2.errors.UniqueViolation:
pass
except psycopg2.Error:
Expand Down
22 changes: 10 additions & 12 deletions src/charm.py
Original file line number Diff line number Diff line change
Expand Up @@ -552,24 +552,22 @@ def enable_disable_extensions(self, database: str = None) -> None:
database: optional database where to enable/disable the extension.
"""
original_status = self.unit.status
extensions = {}
# collect extensions
plugins_exception = {"uuid_ossp": '"uuid-ossp"'}
for plugin in self.config.plugin_keys():
enable = self.config[plugin]

# Enable or disable the plugin/extension.
extension = "_".join(plugin.split("_")[1:-1])
if extension in plugins_exception:
extension = plugins_exception[extension]
self.unit.status = WaitingStatus(
f"{'Enabling' if enable else 'Disabling'} {extension}"
)
try:
self.postgresql.enable_disable_extension(extension, enable, database)
except PostgreSQLEnableDisableExtensionError as e:
logger.exception(
f"failed to {'enable' if enable else 'disable'} {extension} plugin: %s", str(e)
)
self.unit.status = original_status
extension = plugins_exception.get(extension, extension)
extensions[extension] = enable
self.unit.status = WaitingStatus("Updating extensions")
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can no longer tell if we are enabling or disabling and which plugin.

try:
self.postgresql.enable_disable_extensions(extensions, database)
except PostgreSQLEnableDisableExtensionError as e:
logger.exception("failed to change plugins: %s", str(e))
self.unit.status = original_status

def _add_members(self, event) -> None:
"""Add new cluster members.
Expand Down
3 changes: 0 additions & 3 deletions src/relations/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,9 +152,6 @@ def set_up_relation(self, relation: Relation) -> bool:

self.charm.postgresql.create_database(database, user, plugins=plugins)

# Enable/disable extensions in the new database.
self.charm.enable_disable_extensions(database)

Comment on lines -155 to -157
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will be done anyway on line 153

# Build the primary's connection string.
primary = str(
ConnectionString(
Expand Down
10 changes: 0 additions & 10 deletions tests/unit/test_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,14 +183,12 @@ def test_get_extensions(self):
)

@patch("relations.db.DbProvides._update_unit_status")
@patch("charm.PostgresqlOperatorCharm.enable_disable_extensions")
@patch("relations.db.new_password", return_value="test-password")
@patch("relations.db.DbProvides._get_extensions")
def test_set_up_relation(
self,
_get_extensions,
_new_password,
_enable_disable_extensions,
_update_unit_status,
):
with patch.object(PostgresqlOperatorCharm, "postgresql", Mock()) as postgresql_mock:
Expand Down Expand Up @@ -228,7 +226,6 @@ def test_set_up_relation(
self.assertFalse(self.harness.charm.legacy_db_relation.set_up_relation(relation))
postgresql_mock.create_user.assert_not_called()
postgresql_mock.create_database.assert_not_called()
_enable_disable_extensions.assert_not_called()
postgresql_mock.get_postgresql_version.assert_not_called()
_update_unit_status.assert_not_called()

Expand All @@ -244,7 +241,6 @@ def test_set_up_relation(
user = f"relation_id_{self.rel_id}"
postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
postgresql_mock.create_database.assert_called_once_with(DATABASE, user, plugins=[])
_enable_disable_extensions.assert_called_once()
self.assertEqual(postgresql_mock.get_postgresql_version.call_count, 2)
_update_unit_status.assert_called_once()
expected_data = {
Expand All @@ -271,7 +267,6 @@ def test_set_up_relation(
# provided only in the unit databag.
postgresql_mock.create_user.reset_mock()
postgresql_mock.create_database.reset_mock()
_enable_disable_extensions.reset_mock()
postgresql_mock.get_postgresql_version.reset_mock()
_update_unit_status.reset_mock()
with self.harness.hooks_disabled():
Expand All @@ -289,7 +284,6 @@ def test_set_up_relation(
self.assertTrue(self.harness.charm.legacy_db_relation.set_up_relation(relation))
postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
postgresql_mock.create_database.assert_called_once_with(DATABASE, user, plugins=[])
_enable_disable_extensions.assert_called_once()
self.assertEqual(postgresql_mock.get_postgresql_version.call_count, 2)
_update_unit_status.assert_called_once()
self.assertEqual(self.harness.get_relation_data(self.rel_id, self.app), expected_data)
Expand All @@ -299,7 +293,6 @@ def test_set_up_relation(
# Assert that the correct calls were made when the database name is not provided.
postgresql_mock.create_user.reset_mock()
postgresql_mock.create_database.reset_mock()
_enable_disable_extensions.reset_mock()
postgresql_mock.get_postgresql_version.reset_mock()
_update_unit_status.reset_mock()
with self.harness.hooks_disabled():
Expand All @@ -312,7 +305,6 @@ def test_set_up_relation(
self.assertFalse(self.harness.charm.legacy_db_relation.set_up_relation(relation))
postgresql_mock.create_user.assert_not_called()
postgresql_mock.create_database.assert_not_called()
_enable_disable_extensions.assert_not_called()
postgresql_mock.get_postgresql_version.assert_not_called()
_update_unit_status.assert_not_called()
# No data is set in the databags by the database.
Expand All @@ -329,7 +321,6 @@ def test_set_up_relation(
)
self.assertFalse(self.harness.charm.legacy_db_relation.set_up_relation(relation))
postgresql_mock.create_database.assert_not_called()
_enable_disable_extensions.assert_not_called()
postgresql_mock.get_postgresql_version.assert_not_called()
_update_unit_status.assert_not_called()
self.assertIsInstance(self.harness.model.unit.status, BlockedStatus)
Expand All @@ -340,7 +331,6 @@ def test_set_up_relation(
# BlockedStatus due to a PostgreSQLCreateDatabaseError.
self.harness.charm.unit.status = ActiveStatus()
self.assertFalse(self.harness.charm.legacy_db_relation.set_up_relation(relation))
_enable_disable_extensions.assert_not_called()
postgresql_mock.get_postgresql_version.assert_not_called()
_update_unit_status.assert_not_called()
self.assertIsInstance(self.harness.model.unit.status, BlockedStatus)
Expand Down