Ensure no temp leaks

mario/athena.py

@@ -1,10 +1,13 @@
-from pyathena import connect
-from mario.data_extractor import DataExtractor, Configuration, GET_TOTAL_WITHOUT_MEASURE
-import logging
 import datetime
+import logging
 import numbers
+import os
 import re
 
+import awswrangler as wr
+from pyathena import connect
+
+from mario.data_extractor import GET_TOTAL_WITHOUT_MEASURE, Configuration, DataExtractor
 from mario.mapping import rewrite_csv_header_with_fieldmapping
 from mario.options import CsvOptions
 from mario.utils import gzip_file
@@ -95,7 +98,6 @@ def run_query(self) -> str:
         result path in S3
         :return: the S3 path where the result is stored
         """
-        import awswrangler as wr
         import boto3
 
         # Build SQL
@@ -141,10 +143,17 @@ def save_data_as_csv(self, file_path: str, **kwargs):
         wr.s3.download(path=s3_uri, local_file=file_path)
 
         # Rewrite header with FieldMapping
-        rewrite_csv_header_with_fieldmapping(file_path, self.mapping)
+        try:
+            rewrite_csv_header_with_fieldmapping(file_path, self.mapping)
+        except Exception:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+            raise
 
         if options.compress_using_gzip:
             gz_path = gzip_file(file_path)
+            # Delete original
+            os.remove(file_path)
             return gz_path
 
         return file_path

mario/mapping.py

@@ -1,10 +1,9 @@
-from typing import List
-
-import pandas as pd
 import csv
 import os
 import tempfile
 
+import pandas as pd
+
 from mario.utils import to_snake_case
 
 
@@ -13,7 +12,7 @@ class FieldMapping:
     Utility class for holding a set of logical-to-physical item mappings
     """
 
-    def __init__(self, query_format, items: List[str]):
+    def __init__(self, query_format, items: list[str]):
         self._format = query_format
 
         self.as_physical = {}
@@ -48,23 +47,30 @@ def rewrite_csv_header_with_fieldmapping(local_path, field_mapping):
     fd, temp_path = tempfile.mkstemp(dir=dir_name)
     os.close(fd)
 
-    with open(local_path, "r", encoding="utf-8", newline="") as src, \
-         open(temp_path, "w", encoding="utf-8", newline="") as dst:
+    try:
+        with open(local_path, "r", encoding="utf-8", newline="") as src, \
+             open(temp_path, "w", encoding="utf-8", newline="") as dst:
+
+            reader = csv.reader(src)
+            writer = csv.writer(dst, quoting=csv.QUOTE_MINIMAL)
+
+            # --- Read & rewrite only the first line ---
+            physical_header = next(reader)
+            logical_header = [
+                field_mapping.as_logical.get(col, col)
+                for col in physical_header
+            ]
+            writer.writerow(logical_header)
 
-        reader = csv.reader(src)
-        writer = csv.writer(dst, quoting=csv.QUOTE_MINIMAL)
+            # --- Stream the rest unchanged ---
+            for row in reader:
+                writer.writerow(row)
 
-        # --- Read & rewrite only the first line ---
-        physical_header = next(reader)
-        logical_header = [
-            field_mapping.as_logical.get(col, col)
-            for col in physical_header
-        ]
-        writer.writerow(logical_header)
+        # Atomic replace of original file
+        os.replace(temp_path, local_path)
 
-        # --- Stream the rest unchanged ---
-        for row in reader:
-            writer.writerow(row)
+    except Exception:
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+        raise
 
-    # Atomic replace of original file
-    os.replace(temp_path, local_path)

mario/utils.py

@@ -1,6 +1,8 @@
+import logging
 import os
+import shutil
 from datetime import datetime
-import logging
+
 logger = logging.getLogger(__name__)
 
 
@@ -25,11 +27,15 @@ def to_snake_case(name: str) -> str:
 
 def gzip_file(input_path, output_path=None):
     import gzip
-    import shutil
     if output_path is None:
         output_path = input_path + ".gz"
-    with open(input_path, "rb") as f_in, gzip.open(output_path, "wb") as f_out:
-        shutil.copyfileobj(f_in, f_out)
+    try:
+        with open(input_path, "rb") as f_in, gzip.open(output_path, "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+    except Exception:
+        if os.path.exists(output_path):
+            os.remove(output_path)
+        raise
     return output_path