From 0863d317d65f2bd4e623093f00a9411c90786fa6 Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushr2@illinois.edu>
Date: Sat, 31 Aug 2019 23:14:40 +0530
Subject: [PATCH] Add support for stage logs option (#25)

---
 README.md             | 2 ++
 chainlink/__init__.py | 9 +++++----
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index a2c324c..3ea7ccb 100644
--- a/README.md
+++ b/README.md
@@ -43,6 +43,8 @@ stages = [{
   "privileged": True,
   # the number of seconds until the container is killed (optional, defaults to 30)
   "timeout": 30,
+  # enable saving the logs from this stage
+  "logs": False,
   # container environment additions/overrides (optional, defaults to none)
   "env": {
     "VAR1": "1"
diff --git a/chainlink/__init__.py b/chainlink/__init__.py
index 6587283..56c81b3 100644
--- a/chainlink/__init__.py
+++ b/chainlink/__init__.py
@@ -108,11 +108,12 @@ async def _run_stage(self, stage, mount, environ):
         result = {
             "data": self.client.api.inspect_container(container.id)["State"],
             "killed": killed,
-            "logs": {
-                "stdout": container.logs(stderr=False, timestamps=True),
-                "stderr": container.logs(stdout=False, timestamps=True),
-            },
+            "logs": {"stdout": None, "stderr": None},
         }
+        if stage.get("logs", True):
+            result["logs"]["stdout"] = container.logs(stderr=False, timestamps=True)
+            result["logs"]["stderr"] = container.logs(stdout=False, timestamps=True)
+
         result["success"] = (not killed) and (result["data"]["ExitCode"] == 0)
         container.remove()