Support workload identity in flush manager service

ingestion-edge/ingestion_edge/flush_manager.py

@@ -75,6 +75,11 @@
 parser.add_argument(
     "--namespace", default="default", help="Kubernetes namespace to use",
 )
+parser.add_argument(
+    "--service-account-name",
+    default="default",
+    help="Kubernetes service account name to use",
+)
 parser.add_argument(
     "--claim-prefix",
     default="queue-",
@@ -148,6 +153,7 @@ def _create_flush_job(
     image: str,
     name: str,
     namespace: str,
+    service_account_name: str,
 ) -> V1Job:
     logger.info(f"creating job: {name}")
     try:
@@ -182,6 +188,7 @@ def _create_flush_job(
                                     ),
                                 )
                             ],
+                            service_account_name=service_account_name,
                         )
                     )
                 ),
@@ -201,6 +208,7 @@ def flush_released_pvs(
     env: List[V1EnvVar],
     image: str,
     namespace: str,
+    service_account_name: str,
 ):
     """
     Flush persistent volumes.
@@ -224,7 +232,9 @@ def flush_released_pvs(
             if pv.status.phase != "Bound":
                 pvc = _create_pvc(api, name, namespace, pv)
                 _bind_pvc(api, pv, pvc)
-            _create_flush_job(batch_api, command, env, image, name, namespace)
+            _create_flush_job(
+                batch_api, command, env, image, name, namespace, service_account_name
+            )
 
 
 def delete_complete_jobs(api: CoreV1Api, batch_api: BatchV1Api, namespace: str):
@@ -292,12 +302,15 @@ def flush_released_pvs_and_delete_complete_jobs(
     env: List[V1EnvVar],
     image: str,
     namespace: str,
+    service_account_name: str,
 ):
     """Flush released persistent volumes then delete complete jobs.
 
     Run sequentially to avoid race conditions.
     """
-    flush_released_pvs(api, batch_api, command, env, image, namespace)
+    flush_released_pvs(
+        api, batch_api, command, env, image, namespace, service_account_name
+    )
     delete_complete_jobs(api, batch_api, namespace)
 
 
@@ -425,6 +438,7 @@ def main():
             args.env,
             args.image,
             args.namespace,
+            args.service_account_name,
         ),
         partial(delete_detached_pvcs, api, args.namespace, args.claim_prefix),
         partial(delete_unschedulable_pods, api, args.namespace),

ingestion-edge/tests/unit/flush_manager/test_create_job.py

@@ -131,7 +131,9 @@ def create_job(namespace, body):
 
     batch_api.create_namespaced_job.side_effect = create_job
 
-    flush_released_pvs(api, batch_api, "command", "env", "image", "namespace")
+    flush_released_pvs(
+        api, batch_api, "command", "env", "image", "namespace", "service-account-name"
+    )
 
     api.list_persistent_volume.assert_called_once_with()
     batch_api.list_namespaced_job.assert_called_once_with("namespace")
@@ -179,7 +181,15 @@ def create_job(namespace, body):
     batch_api.create_namespaced_job.side_effect = create_job
 
     with pytest.raises(ApiException):
-        flush_released_pvs(api, batch_api, "command", "env", "image", "namespace")
+        flush_released_pvs(
+            api,
+            batch_api,
+            "command",
+            "env",
+            "image",
+            "namespace",
+            "service-account-name",
+        )
 
     api.list_persistent_volume.assert_called_once_with()
     batch_api.list_namespaced_job.called_once_with("namespace")
@@ -207,7 +217,15 @@ def create_pvc(namespace, body):
     api.create_namespaced_persistent_volume_claim.side_effect = create_pvc
 
     with pytest.raises(ApiException):
-        flush_released_pvs(api, batch_api, "command", "env", "image", "namespace")
+        flush_released_pvs(
+            api,
+            batch_api,
+            "command",
+            "env",
+            "image",
+            "namespace",
+            "service-account-name",
+        )
 
     api.list_persistent_volume.assert_called_once_with()
     batch_api.list_namespaced_job.called_once_with("namespace")

ingestion-edge/tests/unit/flush_manager/test_main.py

@@ -28,7 +28,7 @@ def test_run_task():
 
 def test_flush_and_delete(api: MagicMock, batch_api: MagicMock):
     flush_released_pvs_and_delete_complete_jobs(
-        api, batch_api, "command", "env", "image", "namespace"
+        api, batch_api, "command", "env", "image", "namespace", "service-account-name"
     )
     api.list_persistent_volume.assert_called_once_with()
     assert [