Merge pull request #3911 from Yelp/revert-3906-u/krall/ensure_service…

…_account

Revert "Split create_or_find_service_account into a non-mutating method and a mutating method"

paasta_tools/kubernetes/application/controller_wrappers.py

@@ -15,7 +15,6 @@
 from paasta_tools.kubernetes_tools import create_deployment
 from paasta_tools.kubernetes_tools import create_pod_disruption_budget
 from paasta_tools.kubernetes_tools import create_stateful_set
-from paasta_tools.kubernetes_tools import ensure_service_account
 from paasta_tools.kubernetes_tools import KubeClient
 from paasta_tools.kubernetes_tools import KubeDeployment
 from paasta_tools.kubernetes_tools import KubernetesDeploymentConfig
@@ -121,25 +120,6 @@ def update_related_api_objects(self, kube_client: KubeClient) -> None:
         """
         self.ensure_pod_disruption_budget(kube_client, self.soa_config.get_namespace())
 
-    def update_dependency_api_objects(self, kube_client: KubeClient) -> None:
-        """
-        Update related Kubernetes API objects that should be updated before the main object,
-        such as service accounts.
-        :param kube_client:
-        """
-        self.ensure_service_account(kube_client)
-
-    def ensure_service_account(self, kube_client: KubeClient) -> None:
-        """
-        Ensure that the service account for this application exists
-        :param kube_client:
-        """
-        ensure_service_account(
-            iam_role=self.soa_config.get_iam_role(),
-            namespace=self.soa_config.get_namespace(),
-            kube_client=kube_client,
-        )
-
     def delete_pod_disruption_budget(self, kube_client: KubeClient) -> None:
         try:
             kube_client.policy.delete_namespaced_pod_disruption_budget(

paasta_tools/kubernetes_tools.py

@@ -2220,9 +2220,9 @@ def get_pod_template_spec(
             annotations["iam.amazonaws.com/role"] = ""
             iam_role = self.get_iam_role()
             if iam_role:
-                pod_spec_kwargs["service_account_name"] = get_service_account_name(
-                    iam_role
-                )
+                pod_spec_kwargs[
+                    "service_account_name"
+                ] = create_or_find_service_account_name(iam_role, self.get_namespace())
                 if fs_group is None:
                     # We need some reasoable default for group id of a process
                     # running inside the container. Seems like most of such
@@ -4050,9 +4050,12 @@ def get_all_limit_ranges(
 _RE_NORMALIZE_IAM_ROLE = re.compile(r"[^0-9a-zA-Z]+")
 
 
-def get_service_account_name(
+def create_or_find_service_account_name(
     iam_role: str,
+    namespace: str,
     k8s_role: Optional[str] = None,
+    kubeconfig_file: Optional[str] = None,
+    dry_run: bool = False,
 ) -> str:
     # the service account is expected to always be prefixed with paasta- as using the actual namespace
     # potentially wastes a lot of characters (e.g., paasta-nrtsearchservices) that could be used for
@@ -4078,17 +4081,12 @@ def get_service_account_name(
             "Expected at least one of iam_role or k8s_role to be passed in!"
         )
 
-    return sa_name
-
-
-def ensure_service_account(
-    iam_role: str,
-    namespace: str,
-    kube_client: KubeClient,
-    k8s_role: Optional[str] = None,
-) -> None:
-    sa_name = get_service_account_name(iam_role, k8s_role)
+    # if someone is dry-running paasta_setup_tron_namespace or some other tool that
+    # calls this function, we probably don't want to mutate k8s state :)
+    if dry_run:
+        return sa_name
 
+    kube_client = KubeClient(config_file=kubeconfig_file)
     if not any(
         sa.metadata and sa.metadata.name == sa_name
         for sa in get_all_service_accounts(kube_client, namespace)
@@ -4137,6 +4135,8 @@ def ensure_service_account(
                 namespace=namespace, body=role_binding
             )
 
+    return sa_name
+
 
 def mode_to_int(mode: Optional[Union[str, int]]) -> Optional[int]:
     if mode is not None:

paasta_tools/setup_kubernetes_job.py

@@ -281,7 +281,6 @@ def setup_kube_deployments(
                 "paasta_namespace": app.kube_deployment.namespace,
             }
             try:
-                app.update_dependency_api_objects(kube_client)
                 if (
                     app.kube_deployment.service,
                     app.kube_deployment.instance,

paasta_tools/tron_tools.py

@@ -62,7 +62,7 @@
 
 from paasta_tools.kubernetes_tools import (
     allowlist_denylist_to_requirements,
-    get_service_account_name,
+    create_or_find_service_account_name,
     limit_size_with_hash,
     raw_selectors_to_requirements,
     to_node_label,
@@ -354,11 +354,15 @@ def build_spark_config(self) -> Dict[str, str]:
             "spark.kubernetes.executor.label.yelp.com/owner", self.get_team()
         )
 
+        # We need to make sure the Service Account used by the executors has been created.
         # We are using the Service Account created using the provided or default IAM role.
         spark_conf[
             "spark.kubernetes.authenticate.executor.serviceAccountName"
-        ] = get_service_account_name(
+        ] = create_or_find_service_account_name(
             iam_role=self.get_spark_executor_iam_role(),
+            namespace=spark_tools.SPARK_EXECUTOR_NAMESPACE,
+            kubeconfig_file=system_paasta_config.get_spark_kubeconfig(),
+            dry_run=self.for_validation,
         )
 
         return spark_conf
@@ -948,14 +952,20 @@ def format_tron_action_dict(action_config: TronActionConfig):
 
         result["labels"]["yelp.com/owner"] = "compute_infra_platform_experience"
 
+        # create_or_find_service_account_name requires k8s credentials, and we don't
+        # have those available for CI to use (nor do we check these for normal PaaSTA
+        # services, so we're not doing anything "new" by skipping this)
         if (
             action_config.get_iam_role_provider() == "aws"
             and action_config.get_iam_role()
+            and not action_config.for_validation
         ):
             # this service account will be used for normal Tron batches as well as for Spark drivers
-            result["service_account_name"] = get_service_account_name(
+            result["service_account_name"] = create_or_find_service_account_name(
                 iam_role=action_config.get_iam_role(),
+                namespace=EXECUTOR_TYPE_TO_NAMESPACE[executor],
                 k8s_role=None,
+                dry_run=action_config.for_validation,
             )
 
         # service account token volumes for service authentication

tests/test_kubernetes_tools.py

@@ -89,14 +89,14 @@
 from paasta_tools.kubernetes_tools import allowlist_denylist_to_requirements
 from paasta_tools.kubernetes_tools import create_custom_resource
 from paasta_tools.kubernetes_tools import create_deployment
+from paasta_tools.kubernetes_tools import create_or_find_service_account_name
 from paasta_tools.kubernetes_tools import create_pod_disruption_budget
 from paasta_tools.kubernetes_tools import create_secret
 from paasta_tools.kubernetes_tools import create_secret_signature
 from paasta_tools.kubernetes_tools import create_stateful_set
 from paasta_tools.kubernetes_tools import ensure_namespace
 from paasta_tools.kubernetes_tools import ensure_paasta_api_rolebinding
 from paasta_tools.kubernetes_tools import ensure_paasta_namespace_limits
-from paasta_tools.kubernetes_tools import ensure_service_account
 from paasta_tools.kubernetes_tools import filter_nodes_by_blacklist
 from paasta_tools.kubernetes_tools import filter_pods_by_service_instance
 from paasta_tools.kubernetes_tools import force_delete_pods
@@ -117,7 +117,6 @@
 from paasta_tools.kubernetes_tools import get_secret
 from paasta_tools.kubernetes_tools import get_secret_name_from_ref
 from paasta_tools.kubernetes_tools import get_secret_signature
-from paasta_tools.kubernetes_tools import get_service_account_name
 from paasta_tools.kubernetes_tools import InvalidKubernetesConfig
 from paasta_tools.kubernetes_tools import is_node_ready
 from paasta_tools.kubernetes_tools import is_pod_ready
@@ -4458,7 +4457,7 @@ def test_get_pod_management_policy(config_dict, expected_management_policy):
     assert deployment.get_pod_management_policy() == expected_management_policy
 
 
-def test_ensure_service_account_new():
+def test_create_or_find_service_account_name_new():
     iam_role = "arn:aws:iam::000000000000:role/some_role"
     namespace = "test_namespace"
     k8s_role = None
@@ -4478,8 +4477,8 @@ def test_ensure_service_account_new():
         mock_client.core.list_namespaced_service_account.return_value.items = []
         mock_kube_client.return_value = mock_client
 
-        ensure_service_account(
-            iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
+        assert expected_sa_name == create_or_find_service_account_name(
+            iam_role, namespace=namespace, k8s_role=k8s_role
         )
         mock_client.core.create_namespaced_service_account.assert_called_once_with(
             namespace=namespace,
@@ -4495,7 +4494,7 @@ def test_ensure_service_account_new():
         mock_client.rbac.create_namespaced_role_binding.assert_not_called()
 
 
-def test_ensure_service_account_with_k8s_role_new():
+def test_create_or_find_service_account_name_with_k8s_role_new():
     iam_role = "arn:aws:iam::000000000000:role/some_role"
     namespace = "test_namespace"
     k8s_role = "mega-admin"
@@ -4519,8 +4518,8 @@ def test_ensure_service_account_with_k8s_role_new():
         mock_client.rbac.list_namespaced_role_binding.return_value.items = []
         mock_kube_client.return_value = mock_client
 
-        ensure_service_account(
-            iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
+        assert expected_sa_name == create_or_find_service_account_name(
+            iam_role, namespace=namespace, k8s_role=k8s_role
         )
         mock_client.core.create_namespaced_service_account.assert_called_once_with(
             namespace=namespace,
@@ -4556,7 +4555,7 @@ def test_ensure_service_account_with_k8s_role_new():
         )
 
 
-def test_ensure_service_account_existing():
+def test_create_or_find_service_account_name_existing():
     iam_role = "arn:aws:iam::000000000000:role/some_role"
     namespace = "test_namespace"
     k8s_role = "mega-admin"
@@ -4609,14 +4608,14 @@ def test_ensure_service_account_existing():
         ]
         mock_kube_client.return_value = mock_client
 
-        ensure_service_account(
-            iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
+        assert expected_sa_name == create_or_find_service_account_name(
+            iam_role, namespace=namespace, k8s_role=k8s_role
         )
         mock_client.core.create_namespaced_service_account.assert_not_called()
         mock_client.rbac.create_namespaced_role_binding.assert_not_called()
 
 
-def test_ensure_service_account_existing_create_rb_only():
+def test_create_or_find_service_account_name_existing_create_rb_only():
     iam_role = "arn:aws:iam::000000000000:role/some_role"
     namespace = "test_namespace"
     k8s_role = "mega-admin"
@@ -4649,25 +4648,48 @@ def test_ensure_service_account_existing_create_rb_only():
         mock_client.rbac.list_namespaced_role_binding.return_value.items = []
         mock_kube_client.return_value = mock_client
 
-        ensure_service_account(
-            iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
+        assert expected_sa_name == create_or_find_service_account_name(
+            iam_role, namespace=namespace, k8s_role=k8s_role
         )
         mock_client.core.create_namespaced_service_account.assert_not_called()
         assert mock_client.rbac.create_namespaced_role_binding.called is True
 
 
-def test_ensure_service_account_caps():
+def test_create_or_find_service_account_name_caps():
     iam_role = "arn:aws:iam::000000000000:role/Some_Role"
+    namespace = "test_namespace"
     expected_sa_name = "paasta--arn-aws-iam-000000000000-role-some-role"
     with mock.patch(
         "paasta_tools.kubernetes_tools.kube_config.load_kube_config", autospec=True
-    ):
-        assert expected_sa_name == get_service_account_name(
+    ), mock.patch(
+        "paasta_tools.kubernetes_tools.KubeClient",
+        autospec=False,
+    ) as mock_kube_client:
+        mock_client = mock.Mock()
+        mock_client.core = mock.Mock(spec=kube_client.CoreV1Api)
+        mock_client.core.list_namespaced_service_account.return_value = mock.Mock(
+            spec=V1ServiceAccountList
+        )
+        mock_client.core.list_namespaced_service_account.return_value.items = [
+            V1ServiceAccount(
+                kind="ServiceAccount",
+                metadata=V1ObjectMeta(
+                    name=expected_sa_name,
+                    namespace=namespace,
+                    annotations={"eks.amazonaws.com/role-arn": iam_role},
+                ),
+            )
+        ]
+        mock_kube_client.return_value = mock_client
+
+        assert expected_sa_name == create_or_find_service_account_name(
             iam_role,
+            namespace=namespace,
         )
+        mock_client.core.create_namespaced_service_account.assert_not_called()
 
 
-def test_ensure_service_account_caps_with_k8s():
+def test_create_or_find_service_account_name_caps_with_k8s():
     iam_role = "arn:aws:iam::000000000000:role/Some_Role"
     namespace = "test_namespace"
     k8s_role = "mega-admin"
@@ -4720,8 +4742,8 @@ def test_ensure_service_account_caps_with_k8s():
         ]
         mock_kube_client.return_value = mock_client
 
-        ensure_service_account(
-            iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
+        assert expected_sa_name == create_or_find_service_account_name(
+            iam_role, namespace=namespace, k8s_role=k8s_role
         )
         mock_client.core.create_namespaced_service_account.assert_not_called()
         mock_client.rbac.create_namespaced_role_binding.assert_not_called()

tests/test_tron_tools.py

@@ -1426,7 +1426,7 @@ def test_format_tron_action_dict_paasta_k8s(
             autospec=True,
             return_value=False,
         ), mock.patch(
-            "paasta_tools.tron_tools.get_service_account_name",
+            "paasta_tools.tron_tools.create_or_find_service_account_name",
             autospec=True,
             return_value="some--service--account",
         ), mock.patch(