From bb18c44e478a72f03328777ca4d5e25de87e9f4d Mon Sep 17 00:00:00 2001 From: Vinay Sagar Gonabavi Date: Mon, 9 Sep 2024 20:04:31 -0700 Subject: [PATCH] Create service account with iam_role --- paasta_tools/vitesscluster_tools.py | 49 ++++++++++++++++++++++------- tests/test_vitesscluster_tools.py | 5 +++ 2 files changed, 42 insertions(+), 12 deletions(-) diff --git a/paasta_tools/vitesscluster_tools.py b/paasta_tools/vitesscluster_tools.py index 037289034e..a23a876698 100644 --- 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 index 5bd82a17e4..fbf74facac 100644 --- 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,