Reuse connection for pgaudit conf

Author: dragomirp

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -144,24 +144,17 @@ def __init__(
         self.database = database
         self.system_users = system_users if system_users else []
 
-    def _configure_pgaudit(self, enable: bool) -> None:
-        connection = None
-        try:
-            connection = self._connect_to_database()
-            connection.autocommit = True
-            with connection.cursor() as cursor:
-                if enable:
-                    cursor.execute("ALTER SYSTEM SET pgaudit.log = 'ROLE,DDL,MISC,MISC_SET';")
-                    cursor.execute("ALTER SYSTEM SET pgaudit.log_client TO off;")
-                    cursor.execute("ALTER SYSTEM SET pgaudit.log_parameter TO off;")
-                else:
-                    cursor.execute("ALTER SYSTEM RESET pgaudit.log;")
-                    cursor.execute("ALTER SYSTEM RESET pgaudit.log_client;")
-                    cursor.execute("ALTER SYSTEM RESET pgaudit.log_parameter;")
-                cursor.execute("SELECT pg_reload_conf();")
-        finally:
-            if connection is not None:
-                connection.close()
+    def _configure_pgaudit(self, enable: bool, connection) -> None:
+        with connection.cursor() as cursor:
+            if enable:
+                cursor.execute("ALTER SYSTEM SET pgaudit.log = 'ROLE,DDL,MISC,MISC_SET';")
+                cursor.execute("ALTER SYSTEM SET pgaudit.log_client TO off;")
+                cursor.execute("ALTER SYSTEM SET pgaudit.log_parameter TO off;")
+            else:
+                cursor.execute("ALTER SYSTEM RESET pgaudit.log;")
+                cursor.execute("ALTER SYSTEM RESET pgaudit.log_client;")
+                cursor.execute("ALTER SYSTEM RESET pgaudit.log_parameter;")
+            cursor.execute("SELECT pg_reload_conf();")
 
     def _connect_to_database(
         self, database: Optional[str] = None, database_host: Optional[str] = None
@@ -435,7 +428,7 @@ def enable_disable_extensions(
             for extension, enable in extensions.items():
                 ordered_extensions[extension] = enable
 
-            self._configure_pgaudit(False)
+            self._configure_pgaudit(False, connection)
             # Enable/disabled the extension in each database.
             for database in databases:
                 with self._connect_to_database(
@@ -454,8 +447,8 @@ def enable_disable_extensions(
         except psycopg2.Error as e:
             raise PostgreSQLEnableDisableExtensionError() from e
         finally:
-            self._configure_pgaudit(extensions.get("pgaudit", False))
             if connection is not None:
+                self._configure_pgaudit(extensions.get("pgaudit", False), connection)
                 connection.close()
 
     def _generate_database_privileges_statements(

tests/unit/test_postgresql.py

@@ -1,6 +1,6 @@
 # Copyright 2023 Canonical Ltd.
 # See LICENSE file for licensing details.
-from unittest.mock import call, patch
+from unittest.mock import MagicMock, call, patch
 
 import psycopg2
 import pytest
@@ -460,30 +460,26 @@ def test_build_postgresql_parameters(harness):
 
 
 def test_configure_pgaudit(harness):
-    with patch(
-        "charms.postgresql_k8s.v0.postgresql.PostgreSQL._connect_to_database"
-    ) as _connect_to_database:
-        # Test when pgAudit is enabled.
-        execute = (
-            _connect_to_database.return_value.cursor.return_value.__enter__.return_value.execute
-        )
-        harness.charm.postgresql._configure_pgaudit(True)
-        execute.assert_has_calls([
-            call("ALTER SYSTEM SET pgaudit.log = 'ROLE,DDL,MISC,MISC_SET';"),
-            call("ALTER SYSTEM SET pgaudit.log_client TO off;"),
-            call("ALTER SYSTEM SET pgaudit.log_parameter TO off;"),
-            call("SELECT pg_reload_conf();"),
-        ])
-
-        # Test when pgAudit is disabled.
-        execute.reset_mock()
-        harness.charm.postgresql._configure_pgaudit(False)
-        execute.assert_has_calls([
-            call("ALTER SYSTEM RESET pgaudit.log;"),
-            call("ALTER SYSTEM RESET pgaudit.log_client;"),
-            call("ALTER SYSTEM RESET pgaudit.log_parameter;"),
-            call("SELECT pg_reload_conf();"),
-        ])
+    # Test when pgAudit is enabled.
+    mock_conn = MagicMock()
+    execute = mock_conn.cursor.return_value.__enter__.return_value.execute
+    harness.charm.postgresql._configure_pgaudit(True, mock_conn)
+    execute.assert_has_calls([
+        call("ALTER SYSTEM SET pgaudit.log = 'ROLE,DDL,MISC,MISC_SET';"),
+        call("ALTER SYSTEM SET pgaudit.log_client TO off;"),
+        call("ALTER SYSTEM SET pgaudit.log_parameter TO off;"),
+        call("SELECT pg_reload_conf();"),
+    ])
+
+    # Test when pgAudit is disabled.
+    execute.reset_mock()
+    harness.charm.postgresql._configure_pgaudit(False, mock_conn)
+    execute.assert_has_calls([
+        call("ALTER SYSTEM RESET pgaudit.log;"),
+        call("ALTER SYSTEM RESET pgaudit.log_client;"),
+        call("ALTER SYSTEM RESET pgaudit.log_parameter;"),
+        call("SELECT pg_reload_conf();"),
+    ])
 
 
 def test_validate_group_map(harness):