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

…tack_registrations_for_vitess_pods

U/gonabavi/dreimp 10901 add smartstack registrations for vitess pods

paasta_tools/vitesscluster_tools.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import sys
 from typing import Any
@@ -15,6 +16,7 @@
 from paasta_tools.kubernetes_tools import KubernetesDeploymentConfigDict
 from paasta_tools.kubernetes_tools import limit_size_with_hash
 from paasta_tools.kubernetes_tools import sanitised_cr_name
+from paasta_tools.long_running_service_tools import load_service_namespace_config
 from paasta_tools.utils import BranchDictV2
 from paasta_tools.utils import deep_merge_dictionaries
 from paasta_tools.utils import DEFAULT_SOA_DIR
@@ -134,6 +136,7 @@ class GatewayConfigDict(TypedDict, total=False):
     extraLabels: Dict[str, str]
     replicas: int
     resources: Dict[str, Any]
+    annotations: Mapping[str, Any]
 
 
 class CellConfigDict(TypedDict, total=False):
@@ -149,6 +152,7 @@ class VitessDashboardConfigDict(TypedDict, total=False):
     extraLabels: Dict[str, str]
     replicas: int
     resources: Dict[str, Any]
+    annotations: Mapping[str, Any]
 
 
 class VtAdminConfigDict(TypedDict, total=False):
@@ -162,6 +166,7 @@ class VtAdminConfigDict(TypedDict, total=False):
     readOnly: bool
     apiResources: Dict[str, Any]
     webResources: Dict[str, Any]
+    annotations: Mapping[str, Any]
 
 
 class VtTabletDict(TypedDict, total=False):
@@ -182,6 +187,7 @@ class TabletPoolDict(TypedDict, total=False):
     vttablet: VtTabletDict
     externalDatastore: Dict[str, Any]
     dataVolumeClaimTemplate: Dict[str, Any]
+    annotations: Mapping[str, Any]
 
 
 class ShardTemplateDict(TypedDict, total=False):
@@ -229,6 +235,7 @@ def get_cell_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
+    annotations: Mapping[str, Any],
 ) -> CellConfigDict:
     """
     get vtgate config
@@ -264,6 +271,7 @@ def get_cell_config(
                 "requests": requests,
                 "limits": requests,
             },
+            annotations=annotations,
         ),
     )
     return config
@@ -276,6 +284,7 @@ def get_vitess_dashboard_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
+    annotations: Mapping[str, Any],
 ) -> VitessDashboardConfigDict:
     """
     get vtctld config
@@ -303,6 +312,7 @@ def get_vitess_dashboard_config(
             "requests": requests,
             "limits": requests,
         },
+        annotations=annotations,
     )
     return config
 
@@ -313,6 +323,7 @@ def get_vt_admin_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
+    annotations: Mapping[str, Any],
 ) -> VtAdminConfigDict:
     """
     get vtadmin config
@@ -338,6 +349,7 @@ def get_vt_admin_config(
             "requests": requests,
             "limits": requests,
         },
+        annotations=annotations,
     )
     return config
 
@@ -356,6 +368,7 @@ def get_tablet_pool_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
+    annotations: Mapping[str, Any],
 ) -> TabletPoolDict:
     """
     get vttablet config
@@ -478,6 +491,7 @@ def get_tablet_pool_config(
             "resources": {"requests": {"storage": "10Gi"}},
             "storageClassName": "ebs-csi-gp3",
         },
+        annotations=annotations,
     )
     return config
 
@@ -490,6 +504,7 @@ def get_keyspaces_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
+    annotations: Mapping[str, Any],
 ) -> List[KeyspaceConfigDict]:
     """
     get vitess keyspace config
@@ -548,6 +563,7 @@ def get_keyspaces_config(
                         env,
                         labels,
                         node_affinity,
+                        annotations,
                     )
                     for cell in cells
                 ]
@@ -655,6 +671,22 @@ def get_labels(self) -> Dict[str, str]:
             labels["yelp.com/owner"] = "dre_mysql"
         return labels
 
+    def get_annotations(self) -> Mapping[str, Any]:
+        # get required annotations to be added to the formatted resource before creating or updating custom resource
+        service_namespace_config = load_service_namespace_config(
+            service=self.service, namespace=self.get_nerve_namespace()
+        )
+        system_paasta_config = load_system_paasta_config()
+        has_routable_ip = self.has_routable_ip(
+            service_namespace_config, system_paasta_config
+        )
+        annotations: Mapping[str, Any] = {
+            "smartstack_registrations": json.dumps(self.get_registrations()),
+            "paasta.yelp.com/routable_ip": has_routable_ip,
+        }
+
+        return annotations
+
     def get_vitess_node_affinity(self) -> dict:
         # Workaround from https://github.com/kubernetes-client/python/issues/390
         api_client = ApiClient()
@@ -698,10 +730,17 @@ def get_cells(self) -> List[CellConfigDict]:
         formatted_env = self.get_env_variables()
         labels = self.get_labels()
         node_affinity = self.get_vitess_node_affinity()
+        annotations = self.get_annotations()
 
         return [
             get_cell_config(
-                cell, region, vtgate_resources, formatted_env, labels, node_affinity
+                cell,
+                region,
+                vtgate_resources,
+                formatted_env,
+                labels,
+                node_affinity,
+                annotations,
             )
             for cell in cells
         ]
@@ -714,9 +753,16 @@ def get_vitess_dashboard(self) -> VitessDashboardConfigDict:
         formatted_env = self.get_env_variables()
         labels = self.get_labels()
         node_affinity = self.get_vitess_node_affinity()
+        annotations = self.get_annotations()
 
         return get_vitess_dashboard_config(
-            cells, zk_address, vtctld_resources, formatted_env, labels, node_affinity
+            cells,
+            zk_address,
+            vtctld_resources,
+            formatted_env,
+            labels,
+            node_affinity,
+            annotations,
         )
 
     def get_vtadmin(self) -> VtAdminConfigDict:
@@ -726,9 +772,10 @@ def get_vtadmin(self) -> VtAdminConfigDict:
         formatted_env = self.get_env_variables()
         labels = self.get_labels()
         node_affinity = self.get_vitess_node_affinity()
+        annotations = self.get_annotations()
 
         return get_vt_admin_config(
-            cells, vtadmin_resources, formatted_env, labels, node_affinity
+            cells, vtadmin_resources, formatted_env, labels, node_affinity, annotations
         )
 
     def get_keyspaces(self) -> List[KeyspaceConfigDict]:
@@ -740,9 +787,17 @@ def get_keyspaces(self) -> List[KeyspaceConfigDict]:
         formatted_env = self.get_env_variables()
         labels = self.get_labels()
         node_affinity = self.get_vitess_node_affinity()
+        annotations = self.get_annotations()
 
         return get_keyspaces_config(
-            cells, keyspaces, zk_address, region, formatted_env, labels, node_affinity
+            cells,
+            keyspaces,
+            zk_address,
+            region,
+            formatted_env,
+            labels,
+            node_affinity,
+            annotations,
         )
 
     def get_update_strategy(self) -> Dict[str, str]:

tests/test_vitesscluster_tools.py

@@ -77,6 +77,10 @@
                         }
                     }
                 },
+                "annotations": {
+                    "paasta.yelp.com/routable_ip": "false",
+                    "smartstack_registrations": '["fake_service.fake_instance"]',
+                },
                 "extraEnv": [
                     {
                         "name": "VAULT_ADDR",
@@ -169,6 +173,10 @@
                                             }
                                         }
                                     },
+                                    "annotations": {
+                                        "paasta.yelp.com/routable_ip": "false",
+                                        "smartstack_registrations": '["fake_service.fake_instance"]',
+                                    },
                                     "cell": "fake_cell",
                                     "dataVolumeClaimTemplate": {
                                         "accessModes": ["ReadWriteOnce"],
@@ -353,6 +361,10 @@
                                             }
                                         }
                                     },
+                                    "annotations": {
+                                        "paasta.yelp.com/routable_ip": "false",
+                                        "smartstack_registrations": '["fake_service.fake_instance"]',
+                                    },
                                     "cell": "fake_cell",
                                     "dataVolumeClaimTemplate": {
                                         "accessModes": ["ReadWriteOnce"],
@@ -544,6 +556,10 @@
                 }
             }
         },
+        "annotations": {
+            "paasta.yelp.com/routable_ip": "false",
+            "smartstack_registrations": '["fake_service.fake_instance"]',
+        },
         "cells": ["fake_cell"],
         "extraEnv": [
             {
@@ -586,6 +602,10 @@
                 }
             }
         },
+        "annotations": {
+            "paasta.yelp.com/routable_ip": "false",
+            "smartstack_registrations": '["fake_service.fake_instance"]',
+        },
         "apiAddresses": ["http://localhost:15000"],
         "apiResources": {
             "limits": {"cpu": "100m", "memory": "256Mi"},