ft-orchestration store jq output directly on target instead of temporary files

Author: pascal260303

tools/ft-orchestration/src/collector/jq.py

@@ -1,12 +1,11 @@
 import json
 import logging
-from os import path
 from pathlib import Path
 import tempfile
 import time
 import xml.etree.ElementTree as ET
 
-import pandas as pd
+import numpy as np
 from src.collector.interface import (
     CollectorOutputReaderException,
     CollectorOutputReaderInterface,
@@ -36,6 +35,21 @@
     "ipfix:msgLength": "MSG_LENGTH",
 }
 
+CSV_DTYPES = {
+    "BYTES": np.uint64,
+    "PACKETS": np.uint64,
+    "PROTOCOL": np.uint8,
+    "SRC_PORT": np.uint16,
+    "SRC_IP": str,
+    "DST_PORT": np.uint16,
+    "DST_IP": str,
+    "START_TIME": np.uint64,
+    "END_TIME": np.uint64,
+    "EXPORT_TIME": np.uint32,
+    "SEQ_NUMBER": np.uint32,
+    "MSG_LENGTH": np.uint64,
+}
+
 
 class JQ(CollectorOutputReaderInterface):
     CONFIG_FILE = "config.xml"
@@ -275,36 +289,29 @@ def save_csv(self, csv_file: str):
         csv_file: str
             Path to CSV file. Local file, CSV will be downloaded when collector running on remote.
         """
+        header = CSV_HEADER_TO_ANALYZER_HEADER.values()
 
-        rsync = Rsync(self._executor)
-        filename = path.basename(csv_file)
-        tmp_file = path.join(rsync.get_data_directory(), filename)
-        header = ",".join(
-            [f'"{field}"' for field in CSV_HEADER_TO_ANALYZER_HEADER.keys()]
-        )
+        batch_size = 100_000
+        buffer = []
 
         logging.getLogger().info(
             "Preparing CSV output by calling ipfixcol2 + jq command..."
         )
         start = time.time()
-        Tool(f"echo '{header}' > {tmp_file}", executor=self._executor).run()
         # write csv
-        Tool(f"({self._cmd_csv}) >> {tmp_file}", executor=self._executor).run()
-        end = time.time()
-        logging.getLogger().info("CSV output saved in %.2f seconds.", (end - start))
+        process = AsyncTool(self._cmd_csv, executor=self._executor)
+        process.run()
 
-        start = time.time()
-        tmp_dir = tempfile.mkdtemp()
-        flows_file = rsync.pull_path(tmp_file, tmp_dir)
-        end = time.time()
-        df = pd.read_csv(flows_file)
-        # Filter columns
-        df = df[list(CSV_HEADER_TO_ANALYZER_HEADER.keys())]
-        # rename columns
-        df = df.rename(columns=CSV_HEADER_TO_ANALYZER_HEADER)
+        with open(csv_file, mode="w", buffering=1024**2) as f:
+            f.write(",".join(header) + "\n")
+            for line in process.stdout:
+                buffer.append(line + "\n")
+                if len(buffer) >= batch_size:
+                    f.writelines(buffer)
+                    buffer = []
 
-        df.to_csv(csv_file, index=False)
+            if buffer:
+                f.writelines(buffer)
 
-        logging.getLogger().info(
-            "CSV output downloaded in %.2f seconds.", (end - start)
-        )
+        end = time.time()
+        logging.getLogger().info("CSV output saved in %.2f seconds.", (end - start))