From 5104f277b85e3716cdfc37f28be1f8aa594e28d5 Mon Sep 17 00:00:00 2001
From: Sorin Sbarnea <ssbarnea@redhat.com>
Date: Fri, 6 Nov 2020 16:16:13 +0000
Subject: [PATCH] Make file writing atomic

This should render the need to use wrappers obsolete as all
file writing operations are now atomic, assuring that we either
write the entire file or fail.

That is important as we do not want to end-up serving partial files
with the web-server.

Change-Id: I696e2474b557e6b5fea707a198f32cea721cc150
---
 elastic_recheck/cmd/graph.py               | 14 ++++----------
 elastic_recheck/cmd/uncategorized_fails.py | 13 +++++++------
 elastic_recheck/utils.py                   | 16 ++++++++++++++++
 3 files changed, 27 insertions(+), 16 deletions(-)
 create mode 100644 elastic_recheck/utils.py

diff --git a/elastic_recheck/cmd/graph.py b/elastic_recheck/cmd/graph.py
index da1388ad..e5d60705 100755
--- a/elastic_recheck/cmd/graph.py
+++ b/elastic_recheck/cmd/graph.py
@@ -18,7 +18,6 @@ import argparse
 from datetime import datetime
 import json
 import os
-import sys
 
 from launchpadlib import launchpad
 import pyelasticsearch
@@ -43,6 +42,7 @@ import elastic_recheck.elasticRecheck as er
 from elastic_recheck import log as logging
 import elastic_recheck.query_builder as qb
 import elastic_recheck.results as er_results
+from elastic_recheck.utils import atomic_write
 
 STEP = 3600000
 
@@ -237,15 +237,9 @@ def main():
                      key=lambda bug: -(bug['fails24'] * 100000 + bug['fails']))
 
     jsondata['buglist'] = buglist
-    if args.output:
-        out = open(args.output, 'w')
-    else:
-        out = sys.stdout
-
-    try:
-        out.write(json.dumps(jsondata))
-    finally:
-        out.close()
+    atomic_write(
+        args.output,
+        json.dumps(jsondata))
 
 
 if __name__ == "__main__":
diff --git a/elastic_recheck/cmd/uncategorized_fails.py b/elastic_recheck/cmd/uncategorized_fails.py
index f65eaf66..6b8f993b 100755
--- a/elastic_recheck/cmd/uncategorized_fails.py
+++ b/elastic_recheck/cmd/uncategorized_fails.py
@@ -30,6 +30,7 @@ import elastic_recheck.config as er_config
 import elastic_recheck.elasticRecheck as er
 import elastic_recheck.query_builder as qb
 import elastic_recheck.results as er_results
+from elastic_recheck.utils import atomic_write
 
 LOG = logging.getLogger('eruncategorized')
 
@@ -370,14 +371,14 @@ def main():
                             'ALL_FAILS_QUERY might be broken.', group)
             continue
         data = collect_metrics(classifier, fails, config=config)
+        LOG.info(
+            "Using templates from %s for %s group",
+            opts.templatedir, group)
         engine = setup_template_engine(opts.templatedir, group=group)
         html = classifying_rate(fails, data, engine, classifier, config.ls_url)
-        if opts.output:
-            out_dir = opts.output
-        else:
-            out_dir = os.getcwd()
-        with open(os.path.join(out_dir, group + '.html'), "w") as f:
-            f.write(html)
+        atomic_write(
+            os.path.join(opts.output or os.getcwd(), group + '.html'),
+            html)
 
 
 if __name__ == "__main__":
diff --git a/elastic_recheck/utils.py b/elastic_recheck/utils.py
new file mode 100644
index 00000000..38dd746b
--- /dev/null
+++ b/elastic_recheck/utils.py
@@ -0,0 +1,16 @@
+import os
+import sys
+from typing import Optional
+
+
+def atomic_write(filename: Optional[str], data: str) -> None:
+    """Atomically writes a file.
+
+    If filename is not mentioned, it will write to sys.stdout
+    """
+    if filename:
+        with open(f"{filename}.tmp", "w") as f:
+            f.write(data)
+        os.replace(f"{filename}.tmp", filename)
+    else:
+        sys.stdout.write(data)