Query job logs and traces

Author: lukasbindreiter

tilebox-workflows/tests/jobs/test_client.py

@@ -6,8 +6,18 @@
 from tests.tasks_data import jobs
 
 from _tilebox.grpc.error import NotFoundError
+from tilebox.datasets.query.pagination import Pagination
 from tilebox.datasets.query.time_interval import datetime_to_timestamp
-from tilebox.workflows.data import Job, JobState, uuid_message_to_uuid, uuid_to_uuid_message
+from tilebox.workflows.data import (
+    Job,
+    JobState,
+    LogRecord,
+    QueryJobLogsResponse,
+    QueryJobSpansResponse,
+    Span,
+    uuid_message_to_uuid,
+    uuid_to_uuid_message,
+)
 from tilebox.workflows.jobs.client import JobClient
 from tilebox.workflows.jobs.service import JobService
 from tilebox.workflows.observability.tracing import NoopWorkflowTracer
@@ -34,6 +44,90 @@ def execute(self, context: ExecutionContext) -> None:
         _ = context
 
 
+def test_query_logs_paginates() -> None:
+    next_page_start = uuid4()
+    log_records = [
+        LogRecord(
+            time=datetime.now(tz=timezone.utc),
+            severity_number=9,
+            severity_text="INFO",
+            body="first",
+            trace_id=None,
+            span_id=None,
+            attributes={},
+            runner_attributes={},
+        ),
+        LogRecord(
+            time=datetime.now(tz=timezone.utc),
+            severity_number=9,
+            severity_text="INFO",
+            body="second",
+            trace_id=None,
+            span_id=None,
+            attributes={},
+            runner_attributes={},
+        ),
+    ]
+    telemetry_service = MagicMock()
+    telemetry_service.query_job_logs.side_effect = [
+        QueryJobLogsResponse([log_records[0]], Pagination(starting_after=next_page_start)),
+        QueryJobLogsResponse([log_records[1]], Pagination()),
+    ]
+    job_client = JobClient(MagicMock(), telemetry_service, NoopWorkflowTracer())
+    job_id = uuid4()
+
+    assert job_client.query_logs(job_id) == log_records
+    assert [call.args[1].starting_after for call in telemetry_service.query_job_logs.call_args_list] == [
+        None,
+        next_page_start,
+    ]
+
+
+def test_query_spans_paginates() -> None:
+    next_page_start = uuid4()
+    spans = [
+        Span(
+            start_time=datetime.now(tz=timezone.utc),
+            end_time=datetime.now(tz=timezone.utc),
+            trace_id="00" * 16,
+            span_id="00" * 8,
+            parent_span_id=None,
+            name="first",
+            status_code="STATUS_CODE_UNSET",
+            status_message="",
+            attributes={},
+            runner_attributes={},
+            events=[],
+        ),
+        Span(
+            start_time=datetime.now(tz=timezone.utc),
+            end_time=datetime.now(tz=timezone.utc),
+            trace_id="00" * 16,
+            span_id="00" * 8,
+            parent_span_id=None,
+            name="second",
+            status_code="STATUS_CODE_UNSET",
+            status_message="",
+            attributes={},
+            runner_attributes={},
+            events=[],
+        ),
+    ]
+    telemetry_service = MagicMock()
+    telemetry_service.query_job_spans.side_effect = [
+        QueryJobSpansResponse([spans[0]], Pagination(starting_after=next_page_start)),
+        QueryJobSpansResponse([spans[1]], Pagination()),
+    ]
+    job_client = JobClient(MagicMock(), telemetry_service, NoopWorkflowTracer())
+    job_id = uuid4()
+
+    assert job_client.query_spans(job_id) == spans
+    assert [call.args[1].starting_after for call in telemetry_service.query_job_spans.call_args_list] == [
+        None,
+        next_page_start,
+    ]
+
+
 class MockJobService(JobServiceStub):
     """A mock implementation of the gRPC job service, that stores jobs in memory as a dict."""
 
@@ -112,7 +206,7 @@ def __init__(self) -> None:
         super().__init__()
         service = JobService(MagicMock())
         service.service = MockJobService()  # mock the gRPC service
-        self.job_client = JobClient(service, NoopWorkflowTracer())
+        self.job_client = JobClient(service, MagicMock(), NoopWorkflowTracer())
         self.count_total_submitted = 0
 
     queued_jobs: Bundle[Job] = Bundle("queued_jobs")

tilebox-workflows/tests/observability/test_tracing.py

@@ -53,7 +53,7 @@ def test_workflow_tracers_do_not_share_client_span_processors(
     tracers = [tracing.WorkflowTracer(service=None, url="https://api.tilebox.com", token=None) for _ in range(3)]
 
     for index, tracer in enumerate(tracers):
-        with tracer.start_as_current_span(f"span-{index}"):
+        with tracer.span(f"span-{index}"):
             pass
 
     assert [processor.span_names for processor in span_processors] == [["span-0"], ["span-1"], ["span-2"]]
@@ -66,7 +66,7 @@ def test_workflow_tracers_copy_configured_span_processors_once(
     tracers = [tracing.WorkflowTracer(service=None, url="https://api.tilebox.com", token=None) for _ in range(2)]
 
     for index, tracer in enumerate(tracers):
-        with tracer.start_as_current_span(f"span-{index}"):
+        with tracer.span(f"span-{index}"):
             pass
 
     assert [processor.span_names for processor in span_processors] == [

tilebox-workflows/tilebox/workflows/client.py

@@ -14,6 +14,7 @@
 )
 from tilebox.workflows.jobs.client import JobClient
 from tilebox.workflows.jobs.service import JobService
+from tilebox.workflows.jobs.telemetry_service import TelemetryService
 from tilebox.workflows.observability.logging import (
     OTELLoggingHandler,
     StructuredLogger,
@@ -101,7 +102,7 @@ def jobs(self) -> JobClient:
             A client for the jobs service.
         """
 
-        return JobClient(JobService(self._channel), self._tracer)
+        return JobClient(JobService(self._channel), TelemetryService(self._channel), self._tracer)
 
     def runner(
         self,