Add lifecycle hooks for cloudmap (de)registration operations on vtgat…

…e pods
Yelp · Aug 22, 2024 · 93adf70 · 93adf70
1 parent e492f59
commit 93adf70
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 0 deletions.
diff --git a/paasta_tools/vitesscluster_tools.py b/paasta_tools/vitesscluster_tools.py
@@ -134,6 +134,7 @@ class GatewayConfigDict(TypedDict, total=False):
     extraEnv: List[Union[KVEnvVar, KVEnvVarValueFrom]]
     extraFlags: Dict[str, str]
     extraLabels: Dict[str, str]
+    lifecycle: Dict[str, Dict[str, Dict[str, List[str]]]]
     replicas: int
     resources: Dict[str, Any]
     annotations: Mapping[str, Any]
@@ -236,6 +237,7 @@ def get_cell_config(
     labels: Dict[str, str],
     node_affinity: dict,
     annotations: Mapping[str, Any],
+    aws_region: str,
 ) -> CellConfigDict:
     """
     get vtgate config
@@ -256,6 +258,26 @@ def get_cell_config(
     config = CellConfigDict(
         name=cell,
         gateway=GatewayConfigDict(
+            lifecycle={
+                "preStop": {
+                    "exec": {
+                        "command": [
+                            "/bin/sh",
+                            "-c",
+                            f"/cloudmap/scripts/deregister_from_cloudmap.sh vtgate-{cell} {aws_region}",
+                        ]
+                    }
+                },
+                "postStart": {
+                    "exec": {
+                        "command": [
+                            "/bin/sh",
+                            "-c",
+                            f"/cloudmap/scripts/register_to_cloudmap.sh vtgate-{cell} {aws_region}",
+                        ]
+                    }
+                },
+            },
             affinity={"nodeAffinity": node_affinity},
             extraEnv=updated_vtgate_extra_env,
             extraFlags={
@@ -712,6 +734,10 @@ def get_region(self) -> str:
             sys.exit(1)
         return region
 
+    def get_aws_region(self) -> str:
+        region = self.get_region()
+        return f"{region[:2]}-{region[2:6]}-{region[6:7]}"
+
     def get_global_lock_server(self) -> Dict[str, Dict[str, str]]:
         zk_address = self.config_dict.get("zk_address")
         return {
@@ -725,6 +751,7 @@ def get_global_lock_server(self) -> Dict[str, Dict[str, str]]:
     def get_cells(self) -> List[CellConfigDict]:
         cells = self.config_dict.get("cells")
         region = self.get_region()
+        aws_region = self.get_aws_region()
         vtgate_resources = self.config_dict.get("vtgate_resources")
 
         formatted_env = self.get_env_variables()
@@ -741,6 +768,7 @@ def get_cells(self) -> List[CellConfigDict]:
                 labels,
                 node_affinity,
                 annotations,
+                aws_region,
             )
             for cell in cells
         ]

diff --git a/tests/test_vitesscluster_tools.py b/tests/test_vitesscluster_tools.py
@@ -117,6 +117,26 @@
                     "mysql_auth_vault_ttl": "60s",
                 },
                 "extraLabels": {"tablet_type": "fake_keyspaces_migration"},
+                "lifecycle": {
+                    "postStart": {
+                        "exec": {
+                            "command": [
+                                "/bin/sh",
+                                "-c",
+                                "/cloudmap/scripts/register_to_cloudmap.sh vtgate-fake_cell mo-ck_r-e",
+                            ]
+                        }
+                    },
+                    "preStop": {
+                        "exec": {
+                            "command": [
+                                "/bin/sh",
+                                "-c",
+                                "/cloudmap/scripts/deregister_from_cloudmap.sh vtgate-fake_cell mo-ck_r-e",
+                            ]
+                        }
+                    },
+                },
                 "replicas": 1,
                 "resources": {
                     "limits": {"cpu": "100m", "memory": "256Mi"},