From f33bcfaf87eda51866d9eba9e10c7f2ec72a001c Mon Sep 17 00:00:00 2001
From: Evgeny Antyshev <eantyshev@gmail.com>
Date: Thu, 6 Oct 2022 17:39:51 +0300
Subject: [PATCH] refactor

---
 playground/infrastructure/cd_helper.py | 34 +++++++++++---------------
 playground/infrastructure/helper.py    |  1 +
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/playground/infrastructure/cd_helper.py b/playground/infrastructure/cd_helper.py
index b3d060b429f3e..77ef0748366c5 100644
--- a/playground/infrastructure/cd_helper.py
+++ b/playground/infrastructure/cd_helper.py
@@ -78,27 +78,21 @@ async def _get_outputs(self, examples: List[Example]):
         Args:
             examples: beam examples that should be run
         """
+        async def _populate_fields(example: Example):
+            try:
+                example.compile_output = await client.get_compile_output(example.pipeline_id)
+                example.output = await client.get_run_output(example.pipeline_id)
+                example.logs = await client.get_log(example.pipeline_id)
+                if example.sdk in [SDK_JAVA, SDK_PYTHON]:
+                    example.graph = await client.get_graph(example.pipeline_id, example.filepath)
+            except Exception as e:
+                logging.error(example.link)
+                logging.error(example.compile_output)
+                raise RuntimeError(f"error in {example.name}") from e
+
         async with GRPCClient() as client:
             await get_statuses(client,
                 examples)  # run examples code and wait until all are executed
-            tasks = [client.get_run_output(example.pipeline_id) for example in examples]
-            outputs = await asyncio.gather(*tasks)
-
-            tasks = [client.get_log(example.pipeline_id) for example in examples]
-            logs = await asyncio.gather(*tasks)
-
-            if len(examples) > 0 and examples[0].sdk in [SDK_PYTHON, SDK_JAVA]:
-                tasks = [
-                    client.get_graph(example.pipeline_id, example.filepath)
-                    for example in examples
-                ]
-                graphs = await asyncio.gather(*tasks)
-
-                for graph, example in zip(graphs, examples):
-                    example.graph = graph
-
-            for output, example in zip(outputs, examples):
-                example.output = output
+            tasks = [_populate_fields(example) for example in examples]
+            await asyncio.gather(*tasks)
 
-            for log, example in zip(logs, examples):
-                example.logs = log
diff --git a/playground/infrastructure/helper.py b/playground/infrastructure/helper.py
index eec90df7e330c..53c9a98072e90 100644
--- a/playground/infrastructure/helper.py
+++ b/playground/infrastructure/helper.py
@@ -70,6 +70,7 @@ class Example:
     type: PrecompiledObjectType = PRECOMPILED_OBJECT_TYPE_UNSPECIFIED
     pipeline_id: str = ""
     output: str = ""
+    compile_output: str = ""
     graph: str = ""