From 4feeaeee102cbf5e3dada3c05d44e0495ca68f9a Mon Sep 17 00:00:00 2001
From: Dragomir Penev <6687393+dragomirp@users.noreply.github.com>
Date: Tue, 21 May 2024 15:52:24 +0300
Subject: [PATCH] [MISC] Add experimental max_connections config (#472)

* Add experimental max_connections config

* Add max_connections to restart param test

* Typo
---
 config.yaml                     |  4 ++++
 src/charm.py                    | 10 ++++++++--
 src/config.py                   |  1 +
 tests/integration/test_charm.py |  9 ++++++++-
 tests/unit/test_charm.py        |  1 +
 5 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/config.yaml b/config.yaml
index 158ef85875..35617edbab 100644
--- a/config.yaml
+++ b/config.yaml
@@ -402,3 +402,7 @@ options:
       Allowed values are: from 0 to 2000000000.
     type: int
     default: 150000000
+  experimental_max_connections:
+    type: int
+    description: |
+      [EXPERIMENTAL] Force set max_connections.
diff --git a/src/charm.py b/src/charm.py
index 67eb66ae4a..99592539af 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -1468,7 +1468,7 @@ def _restart(self, event: RunWithLock) -> None:
             return
 
         try:
-            for attempt in Retrying(wait=wait_fixed(3), stop_after_delay=stop_after_delay(300)):
+            for attempt in Retrying(wait=wait_fixed(3), stop=stop_after_delay(300)):
                 with attempt:
                     if not self._can_connect_to_postgresql:
                         assert False
@@ -1547,8 +1547,14 @@ def update_config(self, is_creating_backup: bool = False) -> bool:
             logger.warning("Early exit update_config: Cannot connect to Postgresql")
             return False
 
+        # Use config value if set, calculate otherwise
+        if self.config.experimental_max_connections:
+            max_connections = self.config.experimental_max_connections
+        else:
+            max_connections = max(4 * os.cpu_count(), 100)
+
         self._patroni.bulk_update_parameters_controller_by_patroni({
-            "max_connections": max(4 * os.cpu_count(), 100),
+            "max_connections": max_connections,
             "max_prepared_transactions": self.config.memory_max_prepared_transactions,
         })
 
diff --git a/src/config.py b/src/config.py
index 6627a67b32..c8196b46ba 100644
--- a/src/config.py
+++ b/src/config.py
@@ -98,6 +98,7 @@ class CharmConfig(BaseConfigModel):
     vacuum_autovacuum_vacuum_cost_delay: Optional[float]
     vacuum_autovacuum_vacuum_scale_factor: Optional[float]
     vacuum_vacuum_freeze_table_age: Optional[int]
+    experimental_max_connections: Optional[int]
 
     @classmethod
     def keys(cls) -> list[str]:
diff --git a/tests/integration/test_charm.py b/tests/integration/test_charm.py
index a8735a66e1..12ac5c5a46 100644
--- a/tests/integration/test_charm.py
+++ b/tests/integration/test_charm.py
@@ -173,6 +173,7 @@ async def test_postgresql_parameters_change(ops_test: OpsTest) -> None:
         "memory_max_prepared_transactions": "100",
         "memory_shared_buffers": "128",
         "response_lc_monetary": "en_GB.utf8",
+        "experimental_max_connections": "200",
     })
     await ops_test.model.wait_for_idle(apps=[DATABASE_APP_NAME], status="active", idle_period=30)
     any_unit_name = ops_test.model.applications[DATABASE_APP_NAME].units[0].name
@@ -186,7 +187,12 @@ async def test_postgresql_parameters_change(ops_test: OpsTest) -> None:
             with psycopg2.connect(
                 f"dbname='postgres' user='operator' host='{host}' password='{password}' connect_timeout=1"
             ) as connection, connection.cursor() as cursor:
-                settings_names = ["max_prepared_transactions", "shared_buffers", "lc_monetary"]
+                settings_names = [
+                    "max_prepared_transactions",
+                    "shared_buffers",
+                    "lc_monetary",
+                    "max_connections",
+                ]
                 cursor.execute(
                     sql.SQL("SELECT name,setting FROM pg_settings WHERE name IN ({});").format(
                         sql.SQL(", ").join(sql.Placeholder() * len(settings_names))
@@ -200,6 +206,7 @@ async def test_postgresql_parameters_change(ops_test: OpsTest) -> None:
                 assert settings["max_prepared_transactions"] == "100"
                 assert settings["shared_buffers"] == "128"
                 assert settings["lc_monetary"] == "en_GB.utf8"
+                assert settings["max_connections"] == "200"
         finally:
             connection.close()
 
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 415048e6ad..48907eb4d0 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1135,6 +1135,7 @@ def test_restart(harness):
     with (
         patch("charm.Patroni.restart_postgresql") as _restart_postgresql,
         patch("charm.Patroni.are_all_members_ready") as _are_all_members_ready,
+        patch("charm.PostgresqlOperatorCharm._can_connect_to_postgresql", return_value=True),
     ):
         _are_all_members_ready.side_effect = [False, True, True]