Add pagination

Author: ashangit

src/spark_history_mcp/common/datadog.py

@@ -1,14 +1,13 @@
 import logging
 from datetime import datetime, timedelta
 
-from datadog_api_client import Configuration, ThreadedApiClient
+from datadog_api_client import Configuration, ThreadedApiClient, ApiClient
 from datadog_api_client.v2.api.logs_api import LogsApi
-from datadog_api_client.v2.model.log import Log
 from datadog_api_client.v2.model.logs_list_request import LogsListRequest
 from datadog_api_client.v2.model.logs_list_request_page import LogsListRequestPage
-from datadog_api_client.v2.model.logs_list_response import LogsListResponse
 from datadog_api_client.v2.model.logs_query_filter import LogsQueryFilter
 from datadog_api_client.v2.model.logs_sort import LogsSort
+from pydantic import BaseModel, Field
 
 from spark_history_mcp.common.variable import (
     POD_NAMESPACE,
@@ -21,8 +20,17 @@
 DATADOG_SECRET_KEYS = f"k8s/{POD_NAMESPACE}/{POD_SERVICE_ACCOUNT}/datadog"
 
 
+class LogDD(BaseModel):
+    timestamp: datetime = Field(description="Timestamp when the log has been emitted")
+    message: str = Field(description="Log message")
+    status: str = Field(description="Log level")
+    host: str = Field(description="Host where the logs has been emitted")
+    pod_name: str = Field(description="Pod name where the logs has been emitted")
+
+
 class Datadog:
-    LOG_LIMIT = 1000
+    LIMIT_PER_QUERY_LOGS = 1000
+    MAX_RETURN_LOGS = 100000
 
     def __init__(self):
         vault_api = VaultApi()
@@ -42,13 +50,10 @@ def __init__(self):
         self.configuration.enable_retry = True
         self.configuration.max_retries = 5
 
-    # TODO manage pagination
-    # add yield on each page
-    # see pagination on mcp
     def get_logs(
         self, index_names: list[str], query: str, _from: datetime, to: datetime
-    ) -> list[Log]:
-        with ThreadedApiClient(self.configuration) as api_client:
+    ) -> list[LogDD]:
+        with ApiClient(self.configuration) as api_client:
             logs_api_instance = LogsApi(api_client)
             request = LogsListRequest(
                 filter=LogsQueryFilter(
@@ -59,28 +64,33 @@ def get_logs(
                 ),
                 sort=LogsSort.TIMESTAMP_ASCENDING,
                 page=LogsListRequestPage(
-                    limit=self.LOG_LIMIT,
+                    limit=self.LIMIT_PER_QUERY_LOGS,
                 ),
             )
             try:
-                response: LogsListResponse = logs_api_instance.list_logs(
-                    body=request
-                ).get()
-
-                logs = []
-                if response.data:
-                    for log in response.data:
-                        pod_name = next((tag for tag in log.attributes.get("tags", []) if tag.startswith('pod_name:')), None).replace('pod_name:','')
-                        logs.append(
-                            {
-                                "id": log.id,
-                                "timestamp": log.attributes.timestamp,
-                                "message": log.attributes.get("message", ""),
-                                "status": log.attributes.get("status", ""),
-                                "host": log.attributes.get("host", ""),
-                                "pod_name":pod_name,
-                            }
+                logs: list[LogDD] = []
+                # Use list_logs_with_pagination for automatic pagination
+                for log in logs_api_instance.list_logs_with_pagination(body=request):
+                    pod_name = next(
+                        (
+                            tag
+                            for tag in log.attributes.get("tags", [])
+                            if tag.startswith("pod_name:")
+                        ),
+                        None,
+                    ).replace("pod_name:", "")
+                    logs.append(
+                        LogDD(
+                            timestamp=log.attributes.timestamp,
+                            message=log.attributes.get("message", ""),
+                            status=log.attributes.get("status", ""),
+                            host=log.attributes.get("host", ""),
+                            pod_name=pod_name,
                         )
+                    )
+
+                    if len(logs) >= self.MAX_RETURN_LOGS:
+                        break
 
                 return logs
             except Exception as e:

src/spark_history_mcp/tools/tools.py

@@ -4,7 +4,6 @@
 import logging
 from typing import Any, Dict, List, Optional
 
-from datadog_api_client.v2.model.log import Log
 from yoshi_client.domains.data_eng_infra.shared.libs.py.yoshi_client import Job
 
 from spark_history_mcp.core.app import mcp
@@ -22,7 +21,7 @@
     StageStatus,
     TaskMetricDistributions,
 )
-from ..common.datadog import Datadog
+from ..common.datadog import Datadog, LogDD
 from ..common.yoshi import Yoshi
 
 from ..utils.utils import parallel_execute
@@ -1320,14 +1319,15 @@ def get_job_definition(job_id: str) -> Job:
     return Yoshi(DATACENTER).get_job_definition(job_id)
 
 
+# TODO see to add pagination on mcp
 @mcp.tool()
 def get_spark_job_logs(
     job_id: str,
     retry_attempt: int,
     start_time: datetime,
     end_time: Optional[datetime] = None,
     status: Optional[str] = None,
-) -> list[Log]:
+) -> list[LogDD]:
     """
     Get logs from DataDog for a Spark job execution.