Open port

src/charm.py

@@ -10,6 +10,7 @@
 import socket
 from typing import Dict, Optional
 
+import lightkube
 from charms.grafana_k8s.v0.grafana_dashboard import GrafanaDashboardProvider
 from charms.loki_k8s.v0.loki_push_api import LogProxyConsumer
 from charms.pgbouncer_k8s.v0 import pgb
@@ -117,6 +118,57 @@ def __init__(self, *args):
     #  Charm Lifecycle Hooks
     # =======================
 
+    def _patch_port(self) -> None:
+        """Patch Juju-created k8s service.
+
+        The k8s service will be tied to pod-0 so that the service is auto cleaned by
+        k8s when the last pod is scaled down.
+        """
+        if not self.unit.is_leader():
+            return
+
+        try:
+            logger.debug("Patching k8s service")
+            client = lightkube.Client()
+            pod0 = client.get(
+                res=lightkube.resources.core_v1.Pod,
+                name=self.app.name + "-0",
+                namespace=self.model.name,
+            )
+            service = lightkube.resources.core_v1.Service(
+                metadata=lightkube.models.meta_v1.ObjectMeta(
+                    name=self.app.name,
+                    namespace=self.model.name,
+                    ownerReferences=pod0.metadata.ownerReferences,
+                    labels={
+                        "app.kubernetes.io/name": self.app.name,
+                    },
+                ),
+                spec=lightkube.models.core_v1.ServiceSpec(
+                    ports=[
+                        lightkube.models.core_v1.ServicePort(
+                            name="pgbouncer",
+                            port=self.config["listen_port"],
+                            targetPort=self.config["listen_port"],
+                        )
+                    ],
+                    selector={"app.kubernetes.io/name": self.app.name},
+                ),
+            )
+            client.patch(
+                res=lightkube.resources.core_v1.Service,
+                obj=service,
+                name=service.metadata.name,
+                namespace=service.metadata.namespace,
+                force=True,
+                field_manager=self.app.name,
+                patch_type=lightkube.types.PatchType.MERGE,
+            )
+            logger.debug("Patched k8s service")
+        except lightkube.ApiError:
+            logger.exception("Failed to patch k8s service")
+            raise
+
     @property
     def version(self) -> str:
         """Returns the version Pgbouncer."""
@@ -218,6 +270,11 @@ def _on_pgbouncer_pebble_ready(self, event: PebbleReadyEvent) -> None:
         # Update postgres endpoints in config to match the current state of the charm.
         self.update_postgres_endpoints(reload_pgbouncer=True)
 
+        if JujuVersion.from_environ().supports_open_port_on_k8s:
+            self.unit.open_port("tcp", self.config["listen_port"])
+        else:
+            self._patch_port()
+
     def _on_config_changed(self, event: ConfigChangedEvent) -> None:
         """Handle changes in configuration.
 
@@ -247,8 +304,15 @@ def _on_config_changed(self, event: ConfigChangedEvent) -> None:
             # This emits relation-changed events to every client relation, so only do it when
             # necessary
             self.update_client_connection_info(self.config["listen_port"])
+            old_port = config["pgbouncer"]["listen_port"]
             config["pgbouncer"]["listen_port"] = self.config["listen_port"]
 
+            if JujuVersion.from_environ().supports_open_port_on_k8s:
+                self.unit.close_port("tcp", old_port)
+                self.unit.open_port("tcp", self.config["listen_port"])
+            else:
+                self._patch_port()
+
         self.render_pgb_config(config)
         try:
             if self.check_pgb_running():

tests/unit/test_charm.py

@@ -40,6 +40,7 @@ def setUp(self):
 
         self.rel_id = self.harness.model.relations[PEER_RELATION_NAME][0].id
 
+    @patch("charm.PgBouncerK8sCharm._patch_port")
     @patch("charm.PgBouncerK8sCharm.read_pgb_config")
     @patch("charm.PgBouncerK8sCharm.push_file")
     @patch("charm.PgBouncerK8sCharm.update_client_connection_info")
@@ -54,6 +55,7 @@ def test_on_config_changed(
         _update_connection_info,
         _push_file,
         _read_pgb_config,
+        _,
     ):
         self.harness.add_relation(BACKEND_RELATION_NAME, "postgres")
         _read_pgb_config.side_effect = lambda: PgbConfig(DEFAULT_CONFIG)
@@ -104,10 +106,11 @@ def test_pgbouncer_layer(self):
         layer = self.charm._pgbouncer_layer()
         assert len(layer.services) == self.charm._cores + 2
 
+    @patch("charm.PgBouncerK8sCharm._patch_port")
     @patch("charm.PgBouncerK8sCharm.update_status")
     @patch("ops.model.Container.exec")
     @patch("ops.model.Container.make_dir")
-    def test_on_pgbouncer_pebble_ready(self, _mkdir, _exec, _update_status):
+    def test_on_pgbouncer_pebble_ready(self, _mkdir, _exec, _update_status, _):
         _exec.return_value.wait_output.return_value = ("PGB 1.16.1\nOther things", "")
         self.harness.add_relation(BACKEND_RELATION_NAME, "postgres")
         self.harness.set_leader(True)
@@ -144,13 +147,14 @@ def test_on_pgbouncer_pebble_ready_defer_tls(self, _mkdir, get_tls_files):
         self.assertIsInstance(self.harness.model.unit.status, WaitingStatus)
         self.assertEqual(self.harness.model.unit.status.message, "Waiting for certificates")
 
+    @patch("charm.PgBouncerK8sCharm._patch_port")
     @patch("charm.PgBouncerK8sCharm.update_status")
     @patch("ops.model.Container.exec")
     @patch("charm.PgBouncerK8sCharm.push_tls_files_to_workload")
     @patch("charm.PostgreSQLTLS.get_tls_files")
     @patch("ops.model.Container.make_dir")
     def test_on_pgbouncer_pebble_ready_ensure_tls_files(
-        self, _mkdir, get_tls_files, push_tls_files_to_workload, _exec, _update_status
+        self, _mkdir, get_tls_files, push_tls_files_to_workload, _exec, _update_status, _
     ):
         _exec.return_value.wait_output.return_value = ("", "")
         get_tls_files.return_value = ("key", "ca", "cert")
@@ -195,11 +199,12 @@ def test_read_pgb_config(self):
         read_cfg = self.charm.read_pgb_config()
         self.assertEqual(PgbConfig(read_cfg).render(), test_cfg.render())
 
+    @patch("charm.PgBouncerK8sCharm._patch_port")
     @patch("charm.PgBouncerK8sCharm.check_pgb_running")
     @patch("ops.model.Container.exec")
     @patch("ops.model.Container.make_dir")
     @patch("ops.model.Container.restart")
-    def test_reload_pgbouncer(self, _restart, _mkdir, _exec, _check_pgb_running):
+    def test_reload_pgbouncer(self, _restart, _mkdir, _exec, _check_pgb_running, _):
         self.harness.add_relation(BACKEND_RELATION_NAME, "postgres")
         self.harness.set_leader(True)
         # necessary hooks before we can check reloads