Create service account with iam_role

Yelp · Sep 10, 2024 · bb18c44 · bb18c44
1 parent b5aa43d
commit bb18c44
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 12 deletions.
diff --git a/paasta_tools/vitesscluster_tools.py b/paasta_tools/vitesscluster_tools.py
@@ -12,6 +12,8 @@
 import service_configuration_lib
 from kubernetes.client import ApiClient
 
+from paasta_tools.kubernetes_tools import ensure_service_account
+from paasta_tools.kubernetes_tools import KubeClient
 from paasta_tools.kubernetes_tools import KubernetesDeploymentConfig
 from paasta_tools.kubernetes_tools import KubernetesDeploymentConfigDict
 from paasta_tools.kubernetes_tools import limit_size_with_hash
@@ -41,6 +43,17 @@
 GRPC_PORT = "15999"
 
 
+PodAnnotationsDict = TypedDict(
+    "PodAnnotationsDict",
+    {
+        "smartstack_registrations": str,
+        "paasta.yelp.com/routable_ip": str,
+        "iam.amazonaws.com/role": str,
+    },
+    total=False,
+)
+
+
 # Environment variables
 VTCTLD_EXTRA_ENV = {
     "WEB_PORT": WEB_PORT,
@@ -137,7 +150,7 @@ class GatewayConfigDict(TypedDict, total=False):
     lifecycle: Dict[str, Dict[str, Dict[str, List[str]]]]
     replicas: int
     resources: Dict[str, Any]
-    annotations: Mapping[str, Any]
+    annotations: PodAnnotationsDict
 
 
 class CellConfigDict(TypedDict, total=False):
@@ -153,7 +166,7 @@ class VitessDashboardConfigDict(TypedDict, total=False):
     extraLabels: Dict[str, str]
     replicas: int
     resources: Dict[str, Any]
-    annotations: Mapping[str, Any]
+    annotations: PodAnnotationsDict
 
 
 class VtAdminConfigDict(TypedDict, total=False):
@@ -167,7 +180,7 @@ class VtAdminConfigDict(TypedDict, total=False):
     readOnly: bool
     apiResources: Dict[str, Any]
     webResources: Dict[str, Any]
-    annotations: Mapping[str, Any]
+    annotations: PodAnnotationsDict
 
 
 class VtTabletDict(TypedDict, total=False):
@@ -188,7 +201,7 @@ class TabletPoolDict(TypedDict, total=False):
     vttablet: VtTabletDict
     externalDatastore: Dict[str, Any]
     dataVolumeClaimTemplate: Dict[str, Any]
-    annotations: Mapping[str, Any]
+    annotations: PodAnnotationsDict
 
 
 class ShardTemplateDict(TypedDict, total=False):
@@ -236,7 +249,7 @@ def get_cell_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
-    annotations: Mapping[str, Any],
+    annotations: PodAnnotationsDict,
     aws_region: str,
 ) -> CellConfigDict:
     """
@@ -306,7 +319,7 @@ def get_vitess_dashboard_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
-    annotations: Mapping[str, Any],
+    annotations: PodAnnotationsDict,
 ) -> VitessDashboardConfigDict:
     """
     get vtctld config
@@ -345,7 +358,7 @@ def get_vt_admin_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
-    annotations: Mapping[str, Any],
+    annotations: PodAnnotationsDict,
 ) -> VtAdminConfigDict:
     """
     get vtadmin config
@@ -390,7 +403,7 @@ def get_tablet_pool_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
-    annotations: Mapping[str, Any],
+    annotations: PodAnnotationsDict,
 ) -> TabletPoolDict:
     """
     get vttablet config
@@ -526,7 +539,7 @@ def get_keyspaces_config(
     env: List[Union[KVEnvVar, KVEnvVarValueFrom]],
     labels: Dict[str, str],
     node_affinity: dict,
-    annotations: Mapping[str, Any],
+    annotations: PodAnnotationsDict,
 ) -> List[KeyspaceConfigDict]:
     """
     get vitess keyspace config
@@ -696,7 +709,7 @@ def get_labels(self) -> Dict[str, str]:
             labels["yelp.com/owner"] = "dre_mysql"
         return labels
 
-    def get_annotations(self) -> Mapping[str, Any]:
+    def get_annotations(self) -> PodAnnotationsDict:
         # 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()
@@ -705,12 +718,24 @@ def get_annotations(self) -> Mapping[str, Any]:
         has_routable_ip = self.has_routable_ip(
             service_namespace_config, system_paasta_config
         )
-        annotations: Mapping[str, Any] = {
+        annotations: PodAnnotationsDict = {
             "smartstack_registrations": json.dumps(self.get_registrations()),
             "paasta.yelp.com/routable_ip": has_routable_ip,
-            "iam.amazonaws.com/role": self.get_iam_role(),
         }
 
+        if self.get_iam_role_provider() == "aws":
+            annotations["iam.amazonaws.com/role"] = ""
+            iam_role = self.get_iam_role()
+            kube_client = KubeClient()
+            if iam_role:
+                ensure_service_account(
+                    iam_role=iam_role,
+                    namespace=self.get_namespace(),
+                    kube_client=kube_client,
+                )
+        else:
+            annotations["iam.amazonaws.com/role"] = self.get_iam_role()
+
         return annotations
 
     def get_vitess_node_affinity(self) -> dict:

diff --git a/tests/test_vitesscluster_tools.py b/tests/test_vitesscluster_tools.py
@@ -676,7 +676,12 @@ def mock_vitess_deployment_config():
     "paasta_tools.vitesscluster_tools.load_system_paasta_config",
     autospec=True,
 )
+@mock.patch(
+    "paasta_tools.vitesscluster_tools.KubeClient",
+    autospec=True,
+)
 def test_load_vitess_service_instance_configs(
+    mock_kube_client,
     mock_load_system_paasta_config,
     mock_load_vitess_instance_config,
     mock_vitess_deployment_config,