Enforce start time and make datadog client a singleton

Author: ashangit

src/spark_history_mcp/common/datadog.py

@@ -1,7 +1,8 @@
 import logging
-from datetime import datetime, timedelta
+from datetime import datetime
+from threading import Lock
 
-from datadog_api_client import Configuration, ThreadedApiClient, ApiClient
+from datadog_api_client import Configuration, ApiClient
 from datadog_api_client.v2.api.logs_api import LogsApi
 from datadog_api_client.v2.model.logs_list_request import LogsListRequest
 from datadog_api_client.v2.model.logs_list_request_page import LogsListRequestPage
@@ -29,7 +30,42 @@ class LogDD(BaseModel):
     pod_name: str = Field(description="Pod name where the logs has been emitted")
 
 
-class Datadog:
+class SingletonMeta(type):
+    """
+    This is a thread-safe implementation of Singleton.
+    """
+
+    _instances = {}
+
+    _lock: Lock = Lock()
+    """
+    We now have a lock object that will be used to synchronize threads during
+    first access to the Singleton.
+    """
+
+    def __call__(cls, *args, **kwargs):
+        """
+        Possible changes to the value of the `__init__` argument do not affect
+        the returned instance.
+        """
+        # Now, imagine that the program has just been launched. Since there's no
+        # Singleton instance yet, multiple threads can simultaneously pass the
+        # previous conditional and reach this point almost at the same time. The
+        # first of them will acquire lock and will proceed further, while the
+        # rest will wait here.
+        with cls._lock:
+            # The first thread to acquire the lock, reaches this conditional,
+            # goes inside and creates the Singleton instance. Once it leaves the
+            # lock block, a thread that might have been waiting for the lock
+            # release may then enter this section. But since the Singleton field
+            # is already initialized, the thread won't create a new object.
+            if cls not in cls._instances:
+                instance = super().__call__(*args, **kwargs)
+                cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+class Datadog(metaclass=SingletonMeta):
     LIMIT_PER_QUERY_LOGS = 1000
     MAX_RETURN_LOGS = 100000
 
@@ -46,6 +82,7 @@ def __init__(self):
         app_key = vault_api.get_secret_kv_store(DATADOG_SECRET_KEYS, "dd_app_key")
 
         self.configuration = Configuration()
+        self.configuration.server_variables["site"] = "datadoghq.com"
         self.configuration.api_key["apiKeyAuth"] = api_key
         self.configuration.api_key["appKeyAuth"] = app_key
         self.configuration.enable_retry = True

src/spark_history_mcp/tools/tools.py

@@ -1363,8 +1363,8 @@ def get_spark_job_logs(
 @mcp.tool()
 def get_operator_logs(
     job_id: str,
+    start_time: datetime,
     status: Optional[str] = None,
-    start_time: Optional[datetime] = None,
     end_time: Optional[datetime] = None,
 ) -> List[LogDD]:
     """
@@ -1375,8 +1375,8 @@ def get_operator_logs(
 
     Args:
         job_id: The mortar/yoshi job ID to get logs for
+        start_time: start time to get logs from
         status: Optional status to filter logs by (e.g. error, warn, info)
-        start_time: Optional start time to get logs from
         end_time: Optional end time to get logs until (defaults to current time)
 
     Returns:
@@ -1400,8 +1400,8 @@ def get_operator_logs(
 @mcp.tool()
 def get_workflow_logs(
     workflow_id: str,
+    start_time: datetime,
     status: Optional[str] = None,
-    start_time: Optional[datetime] = None,
     end_time: Optional[datetime] = None,
 ) -> List[LogDD]:
     """
@@ -1415,8 +1415,8 @@ def get_workflow_logs(
 
     Args:
         workflow_id: The workflow ID to get logs for
+        start_time: start time to get logs from
         status: Optional status to filter logs by (e.g. error, warn, info)
-        start_time: Optional start time to get logs from
         end_time: Optional end time to get logs until (defaults to current time)
 
     Returns:
@@ -1441,8 +1441,8 @@ def get_workflow_logs(
 def get_admission_logs(
     job_id: str,
     app_id: str,
+    start_time: datetime,
     status: Optional[str] = None,
-    start_time: Optional[datetime] = None,
     end_time: Optional[datetime] = None,
 ) -> List[LogDD]:
     """
@@ -1457,8 +1457,8 @@ def get_admission_logs(
     Args:
         job_id: The job ID to get admission logs for
         app_id: The Spark application ID
+        start_time: start time to get logs from
         status: Optional status to filter logs by (e.g. error, warn, info)
-        start_time: Optional start time to get logs from
         end_time: Optional end time to get logs until (defaults to current time)
 
     Returns: