Merge 1f09f7c into 40d15fd

paasta_tools/kubernetes_tools.py

@@ -19,6 +19,7 @@
 import logging
 import math
 import os
+import re
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
@@ -85,6 +86,7 @@
 from kubernetes.client import V1Secret
 from kubernetes.client import V1SecretKeySelector
 from kubernetes.client import V1SecurityContext
+from kubernetes.client import V1ServiceAccount
 from kubernetes.client import V1StatefulSet
 from kubernetes.client import V1StatefulSetSpec
 from kubernetes.client import V1TCPSocketAction
@@ -1262,7 +1264,6 @@ def get_pod_template_spec(
         )
         annotations: Dict[str, Any] = {
             "smartstack_registrations": json.dumps(self.get_registrations()),
-            "iam.amazonaws.com/role": self.get_iam_role(),
             "paasta.yelp.com/routable_ip": "true"
             if service_namespace_config.is_in_smartstack()
             else "false",
@@ -1322,6 +1323,16 @@ def get_pod_template_spec(
                 "termination_grace_period_seconds"
             ] = termination_grace_period
 
+        if self.get_iam_role_provider() == "aws":
+            annotations["iam.amazonaws.com/role"] = ""
+            iam_role = self.get_iam_role()
+            if iam_role:
+                pod_spec_kwargs[
+                    "serviceAccountName"
+                ] = create_or_find_service_account_name(iam_role)
+        else:
+            annotations["iam.amazonaws.com/role"] = self.get_iam_role()
+
         return V1PodTemplateSpec(
             metadata=V1ObjectMeta(
                 labels={
@@ -2568,3 +2579,31 @@ def to_node_label(label: str) -> str:
     }:
         return f"yelp.com/{label}"
     return label
+
+
+def get_all_service_accounts(
+    kube_client: KubeClient, namespace: str,
+) -> Sequence[V1ServiceAccount]:
+    return kube_client.core.list_namespaced_service_account(namespace=namespace).items
+
+
+_RE_NORMALIZE_IAM_ROLE = re.compile(r"[^0-9a-zA-Z]+")
+
+
+def create_or_find_service_account_name(
+    kube_client: KubeClient, iam_role: str, namespace: str = "paasta"
+) -> str:
+    sa_name = "paasta--" + _RE_NORMALIZE_IAM_ROLE.sub("-", iam_role)
+    if not any(
+        sa.name == sa_name for sa in get_all_service_accounts(kube_client, namespace)
+    ):
+        sa = V1ServiceAccount(
+            kind="ServiceAccount",
+            metadata=V1ObjectMeta(
+                name=sa_name,
+                namespace=namespace,
+                annotations={"eks.amazonaws.com/role-arn": iam_role},
+            ),
+        )
+        kube_client.core.create_namespaced_service_account(namespace=namespace, body=sa)
+    return sa_name

paasta_tools/long_running_service_tools.py

@@ -31,6 +31,7 @@
 class LongRunningServiceConfigDict(InstanceConfigDict, total=False):
     drain_method: str
     iam_role: str
+    iam_role_provider: str
     container_port: int
     drain_method_params: Dict
     healthcheck_cmd: str
@@ -203,6 +204,9 @@ def get_replication_crit_percentage(self) -> int:
     def get_iam_role(self) -> str:
         return self.config_dict.get("iam_role", "")
 
+    def get_iam_role_provider(self) -> str:
+        return self.config_dict.get("iam_role_provider", "kiam")
+
     def get_healthcheck_uri(
         self, service_namespace_config: ServiceNamespaceConfig
     ) -> str: