From 930aae727c823a97022e3d34cd0480b30607fc1a Mon Sep 17 00:00:00 2001
From: Terje Kvernes <terjekv@users.noreply.github.com>
Date: Wed, 21 Feb 2024 13:28:32 +0100
Subject: [PATCH] Use corrolation_id for logging across commands. (#213)

---
 mreg_cli/cli.py           |  5 +++++
 mreg_cli/utilities/api.py | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/mreg_cli/cli.py b/mreg_cli/cli.py
index b4683683..0f03bf2a 100644
--- a/mreg_cli/cli.py
+++ b/mreg_cli/cli.py
@@ -28,6 +28,7 @@
 from mreg_cli.help_formatter import CustomHelpFormatter
 from mreg_cli.outputmanager import OutputManager
 from mreg_cli.types import CommandFunc, Flag
+from mreg_cli.utilities.api import create_and_set_corrolation_id
 from mreg_cli.utilities.api import logout as _force_logout
 
 if TYPE_CHECKING:
@@ -246,6 +247,10 @@ def process_command_line(self, line: str) -> None:
         # Set the command that generated the output
         # Also remove filters and other noise.
         cmd = output.from_command(line)
+        # Create and set the corrolation id, using the cleaned command
+        # as the suffix. This is used to track the command in the logs
+        # on the server side.
+        create_and_set_corrolation_id(cmd)
         # Run the command
         cli.parse(cmd)
         # Render the output
diff --git a/mreg_cli/utilities/api.py b/mreg_cli/utilities/api.py
index 3ee321bf..9566a9fa 100644
--- a/mreg_cli/utilities/api.py
+++ b/mreg_cli/utilities/api.py
@@ -7,9 +7,11 @@
 import json
 import logging
 import os
+import re
 import sys
 from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Optional, Union, cast, overload
 from urllib.parse import urljoin
+from uuid import uuid4
 
 import requests
 
@@ -42,6 +44,22 @@ def error(msg: Union[str, Exception], code: int = os.EX_UNAVAILABLE) -> NoReturn
     sys.exit(code)
 
 
+def create_and_set_corrolation_id(suffix: str) -> str:
+    """Set currently active corrolation id.
+
+    This will take a suffix and append it to a generated UUIDv4 and set it as the corrolation id.
+
+    :param suffix: The suffix to use for the corrolation id.
+
+    :returns: The generated corrolation id.
+    """
+    suffix = re.sub(r"\s+", "_", suffix)
+    correlation_id = f"{uuid4()}-{suffix}"
+
+    session.headers.update({"X-Correlation-ID": correlation_id})
+    return correlation_id
+
+
 def set_file_permissions(f: str, mode: int) -> None:
     """Set file permissions on a file."""
     try: