Merge pull request #3934 from Yelp/u/gonabavi/u/gonabavi/DREIMP-10901…

…_add_cloudmap_scripts_add_lifecycle_hooks_for_vtgate_pods Add lifecycle hooks for cloudmap (de)registration operations on vtgate pods
Yelp · Sep 6, 2024 · bf2f8cf · bf2f8cf
2 parents c47f72b + 43758bc
commit bf2f8cf
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 2 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={
@@ -523,7 +545,10 @@ def get_keyspaces_config(
         tablet_types = load_system_paasta_config().get_vitess_tablet_types()
         for tablet_type in tablet_types:
             ecosystem = region.split("-")[-1]
-            host = f"mysql-{cluster}-{tablet_type}.dre-{ecosystem}"
+            if tablet_type == "primary":
+                host = f"mysql-{cluster}.dre-{ecosystem}"
+            else:
+                host = f"mysql-{cluster}-{tablet_type}.dre-{ecosystem}"
 
             # We use migration_replication delay for migration tablets and read_replication_delay for everything else
             # Also throttling threshold for refresh and sanitized primaries is set at 30 seconds and everything else at 3 seconds
@@ -712,6 +737,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 +754,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 +771,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"},
@@ -189,7 +209,7 @@
                                             "volumeName": "vttablet-fake-credentials",
                                         },
                                         "database": "fake_keyspaces",
-                                        "host": "mysql-fake_cluster-primary.dre-devc",
+                                        "host": "mysql-fake_cluster.dre-devc",
                                         "port": 3306,
                                         "user": "vt_app",
                                     },