CM-34777 - Add correlation ID (#220)

Author: MarshalX

cycode/cli/printers/printer_base.py

@@ -5,6 +5,7 @@
 import click
 
 from cycode.cli.models import CliError, CliResult
+from cycode.cyclient.headers import get_correlation_id
 
 if TYPE_CHECKING:
     from cycode.cli.models import LocalScanResult
@@ -46,3 +47,6 @@ def print_exception(self, e: Optional[BaseException] = None) -> None:
             message = f'Error: {traceback_message}'
 
         click.secho(message, err=True, fg=self.RED_COLOR_NAME)
+
+        correlation_message = f'Correlation ID: {get_correlation_id()}'
+        click.secho(correlation_message, err=True, fg=self.RED_COLOR_NAME)

cycode/cyclient/config.py

@@ -1,12 +1,12 @@
 import logging
 import os
 import sys
-from typing import Optional
+from typing import Optional, Union
 from urllib.parse import urlparse
 
 from cycode.cli import consts
 from cycode.cli.user_settings.configuration_manager import ConfigurationManager
-from cycode.cyclient.config_dev import DEV_MODE_ENV_VAR_NAME, DEV_TENANT_ID_ENV_VAR_NAME
+from cycode.cyclient import config_dev
 
 
 def _set_io_encodings() -> None:
@@ -37,20 +37,22 @@ def _set_io_encodings() -> None:
 DEFAULT_CONFIGURATION = {
     consts.TIMEOUT_ENV_VAR_NAME: 300,
     consts.LOGGING_LEVEL_ENV_VAR_NAME: logging.INFO,
-    DEV_MODE_ENV_VAR_NAME: 'False',
+    config_dev.DEV_MODE_ENV_VAR_NAME: 'false',
 }
 
 configuration = dict(DEFAULT_CONFIGURATION, **os.environ)
 
 _CREATED_LOGGERS = set()
 
 
-def get_logger(logger_name: Optional[str] = None) -> logging.Logger:
-    config_level = _get_val_as_string(consts.LOGGING_LEVEL_ENV_VAR_NAME)
-    level = logging.getLevelName(config_level)
+def get_logger_level() -> Optional[Union[int, str]]:
+    config_level = get_val_as_string(consts.LOGGING_LEVEL_ENV_VAR_NAME)
+    return logging.getLevelName(config_level)
+
 
+def get_logger(logger_name: Optional[str] = None) -> logging.Logger:
     new_logger = logging.getLogger(logger_name)
-    new_logger.setLevel(level)
+    new_logger.setLevel(get_logger_level())
 
     _CREATED_LOGGERS.add(new_logger)
 
@@ -62,24 +64,24 @@ def set_logging_level(level: int) -> None:
         created_logger.setLevel(level)
 
 
-def _get_val_as_string(key: str) -> str:
+def get_val_as_string(key: str) -> str:
     return configuration.get(key)
 
 
-def _get_val_as_bool(key: str, default: str = '') -> bool:
+def get_val_as_bool(key: str, default: str = '') -> bool:
     val = configuration.get(key, default)
-    return val.lower() in ('true', '1')
+    return val.lower() in {'true', '1'}
 
 
-def _get_val_as_int(key: str) -> Optional[int]:
+def get_val_as_int(key: str) -> Optional[int]:
     val = configuration.get(key)
     if val:
         return int(val)
 
     return None
 
 
-def _is_valid_url(url: str) -> bool:
+def is_valid_url(url: str) -> bool:
     try:
         urlparse(url)
         return True
@@ -92,12 +94,12 @@ def _is_valid_url(url: str) -> bool:
 configuration_manager = ConfigurationManager()
 
 cycode_api_url = configuration_manager.get_cycode_api_url()
-if not _is_valid_url(cycode_api_url):
+if not is_valid_url(cycode_api_url):
     cycode_api_url = consts.DEFAULT_CYCODE_API_URL
 
-timeout = _get_val_as_int(consts.CYCODE_CLI_REQUEST_TIMEOUT_ENV_VAR_NAME)
+timeout = get_val_as_int(consts.CYCODE_CLI_REQUEST_TIMEOUT_ENV_VAR_NAME)
 if not timeout:
-    timeout = _get_val_as_int(consts.TIMEOUT_ENV_VAR_NAME)
+    timeout = get_val_as_int(consts.TIMEOUT_ENV_VAR_NAME)
 
-dev_mode = _get_val_as_bool(DEV_MODE_ENV_VAR_NAME)
-dev_tenant_id = _get_val_as_string(DEV_TENANT_ID_ENV_VAR_NAME)
+dev_mode = get_val_as_bool(config_dev.DEV_MODE_ENV_VAR_NAME)
+dev_tenant_id = get_val_as_string(config_dev.DEV_TENANT_ID_ENV_VAR_NAME)

cycode/cyclient/cycode_client_base.py

@@ -1,33 +1,17 @@
-import platform
 from typing import ClassVar, Dict, Optional
 
 from requests import Response, exceptions, request
 
-from cycode import __version__
 from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError, NetworkError
-from cycode.cli.user_settings.configuration_manager import ConfigurationManager
 from cycode.cyclient import config, logger
-
-
-def get_cli_user_agent() -> str:
-    """Return base User-Agent of CLI.
-
-    Example: CycodeCLI/0.2.3 (OS: Darwin; Arch: arm64; Python: 3.8.16; InstallID: *uuid4*)
-    """
-    app_name = 'CycodeCLI'
-    version = __version__
-
-    os = platform.system()
-    arch = platform.machine()
-    python_version = platform.python_version()
-
-    install_id = ConfigurationManager().get_or_create_installation_id()
-
-    return f'{app_name}/{version} (OS: {os}; Arch: {arch}; Python: {python_version}; InstallID: {install_id})'
+from cycode.cyclient.headers import get_cli_user_agent, get_correlation_id
 
 
 class CycodeClientBase:
-    MANDATORY_HEADERS: ClassVar[Dict[str, str]] = {'User-Agent': get_cli_user_agent()}
+    MANDATORY_HEADERS: ClassVar[Dict[str, str]] = {
+        'User-Agent': get_cli_user_agent(),
+        'X-Correlation-Id': get_correlation_id(),
+    }
 
     def __init__(self, api_url: str) -> None:
         self.timeout = config.timeout

cycode/cyclient/headers.py

@@ -0,0 +1,46 @@
+import platform
+from typing import Optional
+from uuid import uuid4
+
+from cycode import __version__
+from cycode.cli.user_settings.configuration_manager import ConfigurationManager
+from cycode.cyclient import logger
+
+
+def get_cli_user_agent() -> str:
+    """Return base User-Agent of CLI.
+
+    Example: CycodeCLI/0.2.3 (OS: Darwin; Arch: arm64; Python: 3.8.16; InstallID: *uuid4*)
+    """
+    app_name = 'CycodeCLI'
+    version = __version__
+
+    os = platform.system()
+    arch = platform.machine()
+    python_version = platform.python_version()
+
+    install_id = ConfigurationManager().get_or_create_installation_id()
+
+    return f'{app_name}/{version} (OS: {os}; Arch: {arch}; Python: {python_version}; InstallID: {install_id})'
+
+
+class _CorrelationId:
+    _id: Optional[str] = None
+
+    def get_correlation_id(self) -> str:
+        """Get correlation ID.
+
+        Notes:
+            Used across all requests to correlate logs and metrics.
+            It doesn't depend on client instances.
+            Lifetime is the same as the process.
+        """
+        if self._id is None:
+            # example: 16fd2706-8baf-433b-82eb-8c7fada847da
+            self._id = str(uuid4())
+            logger.debug(f'Correlation ID: {self._id}')
+
+        return self._id
+
+
+get_correlation_id = _CorrelationId().get_correlation_id

tests/cli/exceptions/test_handle_scan_errors.py

@@ -59,7 +59,7 @@ def test_handle_exception_verbose(monkeypatch: 'MonkeyPatch') -> None:
     ctx = click.Context(click.Command('path'), obj={'verbose': True, 'output': 'text'})
 
     def mock_secho(msg: str, *_, **__) -> None:
-        assert 'Error:' in msg
+        assert 'Error:' in msg or 'Correlation ID:' in msg
 
     monkeypatch.setattr(click, 'secho', mock_secho)
 

tests/cyclient/test_client_base.py

@@ -1,10 +1,12 @@
 from cycode.cyclient import config
-from cycode.cyclient.cycode_client_base import CycodeClientBase, get_cli_user_agent
+from cycode.cyclient.cycode_client_base import CycodeClientBase
+from cycode.cyclient.headers import get_cli_user_agent, get_correlation_id
 
 
 def test_mandatory_headers() -> None:
     expected_headers = {
         'User-Agent': get_cli_user_agent(),
+        'X-Correlation-Id': get_correlation_id(),
     }
 
     client = CycodeClientBase(config.cycode_api_url)