datashare-python 0.7.0__py3-none-any.whl → 0.7.1__py3-none-any.whl

datashare_python/conftest.py

@@ -8,6 +8,7 @@ import nest_asyncio
 import pytest
 from elasticsearch._async.helpers import async_streaming_bulk
 from icij_common.es import DOC_ROOT_ID, ES_DOCUMENT_TYPE, ID, ESClient
+from icij_common.test_utils import reset_env  # noqa: F401
 from temporalio import workflow
 
 from datashare_python.config import (

datashare_python/interceptors.py

@@ -0,0 +1,198 @@
+import secrets
+from collections.abc import Generator, Mapping
+from contextlib import contextmanager
+from contextvars import ContextVar
+from copy import deepcopy
+from typing import Annotated, Any, NoReturn, Self, TypeVar
+
+from nexusrpc import InputT, OutputT
+from pydantic import Field
+from temporalio.api.common.v1 import Payload
+from temporalio.converter import DataConverter
+from temporalio.worker import (
+    ActivityInboundInterceptor,
+    ContinueAsNewInput,
+    ExecuteActivityInput,
+    ExecuteWorkflowInput,
+    HandleQueryInput,
+    HandleSignalInput,
+    Interceptor,
+    SignalChildWorkflowInput,
+    SignalExternalWorkflowInput,
+    StartActivityInput,
+    StartChildWorkflowInput,
+    StartLocalActivityInput,
+    StartNexusOperationInput,
+    WorkflowInboundInterceptor,
+    WorkflowInterceptorClassInput,
+    WorkflowOutboundInterceptor,
+)
+from temporalio.workflow import (
+    ActivityHandle,
+    ChildWorkflowHandle,
+    NexusOperationHandle,
+)
+
+from .objects import BaseModel
+
+_TRACEPARENT = "traceparent"
+_DEFAULT_PAYLOAD_CONVERTER = DataConverter.default.payload_converter
+
+
+class TraceContext(BaseModel):
+    # https://www.w3.org/TR/trace-context/
+    version: Annotated[str, Field(frozen=True)] = "00"
+    trace_id: str
+    parent_id: str
+    sampled: bool = True
+
+    def __hash__(self) -> int:
+        return hash((self.trace_id, self.parent_id, self.sampled))
+
+    @classmethod
+    def next_span(cls, parent: Self | None) -> Self:
+        new_span_id = secrets.token_hex(8)
+        if parent is None:
+            trace_id = secrets.token_hex(16)
+            return TraceContext(trace_id=trace_id, parent_id=new_span_id)
+        return TraceContext(
+            trace_id=parent.trace_id, parent_id=new_span_id, sampled=parent.sampled
+        )
+
+    @property
+    def traceparent(self) -> str:
+        flags = "01" if self.sampled else "00"
+        return f"{self.version}-{self.trace_id}-{self.parent_id}-{flags}"
+
+    @classmethod
+    def from_traceparent(cls, traceparent: str) -> Self:
+        split = traceparent.split("-")
+        if len(split) != 4:
+            raise ValueError(f"invalid trace parent: {traceparent}")
+        version, trace_id, parent_id, flags = split
+        if version != "00":
+            msg = (
+                f"unsupported trace parent version {version} "
+                f"for traceparent {traceparent}"
+            )
+            raise ValueError(msg)
+        sampled = flags == "01"
+        return cls(trace_id=trace_id, parent_id=parent_id, sampled=sampled)
+
+
+_TRACE_CONTEXT: ContextVar[TraceContext | None] = ContextVar(
+    "trace_context", default=None
+)
+
+
+class TraceContextInterceptor(Interceptor):
+    def workflow_interceptor_class(
+        self,
+        input: WorkflowInterceptorClassInput,  # noqa: A002, ARG002
+    ) -> type[WorkflowInboundInterceptor] | None:
+        return _TraceContextWorkflowInboundInterceptor
+
+    def intercept_activity(
+        self,
+        next: ActivityInboundInterceptor,  # noqa: A002
+    ) -> ActivityInboundInterceptor:
+        return _TraceContextActivityInboundInterceptor(next)
+
+
+class _TraceContextWorkflowInboundInterceptor(WorkflowInboundInterceptor):
+    def init(self, outbound: WorkflowOutboundInterceptor) -> None:
+        with_outbound_trace_ctx = _TraceContextWorkflowOutboundInterceptor(outbound)
+        super().init(with_outbound_trace_ctx)
+
+    async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:  # noqa: A002
+        with _trace_context(input.headers):
+            return await super().execute_workflow(input)
+
+    async def handle_signal(self, input: HandleSignalInput) -> None:  # noqa: A002
+        with _trace_context(input.headers):
+            return await super().handle_signal(input)
+
+    async def handle_query(self, input: HandleQueryInput) -> Any:  # noqa: A002
+        with _trace_context(input.headers):
+            return await super().handle_query(input)
+
+
+class _TraceContextWorkflowOutboundInterceptor(WorkflowOutboundInterceptor):
+    def continue_as_new(self, input: ContinueAsNewInput) -> NoReturn:  # noqa: A002
+        super().continue_as_new(_with_trace_context_header(input))
+
+    async def signal_child_workflow(self, input: SignalChildWorkflowInput) -> None:  # noqa: A002
+        return await super().signal_child_workflow(_with_trace_context_header(input))
+
+    async def signal_external_workflow(
+        self,
+        input: SignalExternalWorkflowInput,  # noqa: A002
+    ) -> None:
+        return await super().signal_external_workflow(_with_trace_context_header(input))
+
+    def start_activity(self, input: StartActivityInput) -> ActivityHandle[Any]:  # noqa: A002
+        return super().start_activity(_with_trace_context_header(input))
+
+    async def start_child_workflow(
+        self,
+        input: StartChildWorkflowInput,  # noqa: A002
+    ) -> ChildWorkflowHandle[Any, Any]:
+        return await super().start_child_workflow(_with_trace_context_header(input))
+
+    def start_local_activity(
+        self,
+        input: StartLocalActivityInput,  # noqa: A002
+    ) -> ActivityHandle[Any]:
+        return super().start_local_activity(_with_trace_context_header(input))
+
+    async def start_nexus_operation(
+        self,
+        input: StartNexusOperationInput[InputT, OutputT],  # noqa: A002
+    ) -> NexusOperationHandle[OutputT]:
+        return await super().start_nexus_operation(_with_trace_context_header(input))
+
+
+class _TraceContextActivityInboundInterceptor(ActivityInboundInterceptor):
+    async def execute_activity(self, input: ExecuteActivityInput) -> Any:  # noqa: A002
+        with _trace_context(input.headers):
+            return await super().execute_activity(input)
+
+
+def get_trace_context() -> TraceContext | None:
+    return _TRACE_CONTEXT.get()
+
+
+@contextmanager
+def _trace_context(headers: Mapping[str, Payload]) -> Generator[None, None, None]:
+    ctx = headers.get(_TRACEPARENT)
+    if ctx is not None:
+        ctx = _DEFAULT_PAYLOAD_CONVERTER.from_payloads(
+            [headers.get(_TRACEPARENT)], None
+        )[0]
+        ctx = TraceContext.from_traceparent(ctx)
+    else:
+        ctx = TraceContext.next_span(None)
+    tok = None
+    try:
+        tok = _TRACE_CONTEXT.set(ctx)
+        yield
+    finally:
+        if tok is not None:
+            _TRACE_CONTEXT.reset(tok)
+
+
+InputWithHeaders = TypeVar("InputWithHeaders")
+
+
+def _with_trace_context_header(
+    input_with_headers: InputWithHeaders,
+) -> InputWithHeaders:
+    ctx = get_trace_context()
+    if ctx is None:
+        return input_with_headers
+    new_obj = deepcopy(input_with_headers)
+    next_ctx = TraceContext.next_span(ctx)
+    new_obj.headers[_TRACEPARENT] = _DEFAULT_PAYLOAD_CONVERTER.to_payload(
+        next_ctx.traceparent
+    )
+    return new_obj

datashare_python/logging_.py

@@ -12,20 +12,17 @@ from pythonjsonlogger.orjson import OrjsonFormatter
 from temporalio import activity, workflow
 
 from .config import LogLevel
+from .interceptors import get_trace_context
 
-_ACT_LOGGER_ATTRS = [
-    "activity_type",
-    "activity_id",
-    "activity_run_id",
-]
-
-_WF_LOGGED_ATTRS = [
-    "workflow_type",
-    "workflow_id",
-    "workflow_run_id",
-]
+_ACT_LOGGER_ATTRS = ["activity_type", "activity_id", "activity_run_id"]
+_WF_LOGGED_ATTRS = ["workflow_type", "workflow_id", "workflow_run_id"]
+_TRACE_CONTEXT_ATTRS = ["trace_id", "parent_id", "traceparent"]
 _LOGGED_ATTRIBUTES = (
-    copy(RESERVED_ATTRS) + _WF_LOGGED_ATTRS + _ACT_LOGGER_ATTRS + ["worker_id"]
+    copy(RESERVED_ATTRS)
+    + _WF_LOGGED_ATTRS
+    + _ACT_LOGGER_ATTRS
+    + _TRACE_CONTEXT_ATTRS
+    + ["worker_id"]
 )
 
 
@@ -48,7 +45,7 @@ def _get_worker_handlers(
 ) -> list[logging.Handler]:
     stream_handler = logging.StreamHandler(sys.stderr)
     if in_json:
-        fmt = _json_formatter(datefmt=DATE_FMT, worker_id=worker_id)
+        fmt = _json_formatter(datefmt=DATE_FMT)
     else:
         if worker_id is not None:
             fmt = STREAM_HANDLER_FMT_WITH_WORKER_ID
@@ -75,13 +72,15 @@ class WorkerFilter(logging.Filter):
             act_info = activity.info()
             for attr in _ACT_LOGGER_ATTRS:
                 setattr(record, attr, getattr(act_info, attr))
+        trace_context = get_trace_context()
+        if trace_context is not None:
+            for attr in _TRACE_CONTEXT_ATTRS:
+                setattr(record, attr, getattr(trace_context, attr))
         return True
 
 
-def _json_formatter(datefmt: str, worker_id: str) -> BaseJsonFormatter:
+def _json_formatter(datefmt: str) -> BaseJsonFormatter:
     fmt = OrjsonFormatter(  # let's keep logging as fast as possible
-        _LOGGED_ATTRIBUTES,
-        extra={"worker_id": worker_id},
-        datefmt=datefmt,
+        _LOGGED_ATTRIBUTES, datefmt=datefmt
     )
     return fmt

datashare_python/utils.py

@@ -338,6 +338,8 @@ def activity_defn(
     retriables: set[type[Exception]] = None,
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
     def decorator(activity_fn: Callable[P, T]) -> Callable[P, T]:
+        # TODO: some of these could probably be reimplemented more elegantly using
+        #  temporal interceptors: https://docs.temporal.io/develop/python/workers/interceptors
         activity_fn = positional_args_only(activity_fn)
         activity_fn = with_retriables(retriables)(activity_fn)
         if supports_progress(activity_fn):

datashare_python/worker.py

@@ -16,6 +16,7 @@ from temporalio.worker import PollerBehaviorSimpleMaximum, Worker
 from .config import WorkerConfig
 from .dependencies import with_dependencies
 from .discovery import Activity
+from .interceptors import TraceContextInterceptor
 from .types_ import ContextManagerFactory, TemporalClient
 
 logger = logging.getLogger(__name__)
@@ -84,9 +85,10 @@ def datashare_worker(
         max_concurrent_activities = 1
         if workflows:
             logger.warning(_SEPARATE_IO_AND_CPU_WORKERS)
-
+    interceptors = [TraceContextInterceptor()]
     return DatashareWorker(
         client,
+        interceptors=interceptors,
         identity=worker_id,
         workflows=workflows,
         activities=activities,

{datashare_python-0.7.0.dist-info → datashare_python-0.7.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.7.0
+Version: 0.7.1
 Summary: Manage Pythoœn tasks and local resources in Datashare
 Project-URL: Homepage, https://icij.github.io/datashare-python/
 Project-URL: Documentation, https://icij.github.io/datashare-python/

{datashare_python-0.7.0.dist-info → datashare_python-0.7.1.dist-info}/RECORD

@@ -2,25 +2,26 @@ datashare_python/.gitignore,sha256=e-SRgnvGGdsjRrqgKsTzALz6Obx8IYiOjr0yaAxT6v8,2
 datashare_python/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datashare_python/__main__.py,sha256=g-fvS46zl9umKmGrSpl-OG-8PSuZgjqvTCqjpsZtSps,101
 datashare_python/config.py,sha256=Q4iu3ZGaQB7npaHJWclxPIfgzZTf_-8VxxhXrB9nlpE,3928
-datashare_python/conftest.py,sha256=_VESUGax8wt0OlmdJmDFyTJlGO14bt1ccXsjHor0XVk,8592
+datashare_python/conftest.py,sha256=MrmQKFcUipm_qn-cHsLovZMwMMtVxyK0s1lmKEx54bc,8651
 datashare_python/constants.py,sha256=a8-ceZKBVMXydcoNQ35fSjFjxeJ7dt-N6eAvqtPpf9g,320
 datashare_python/dependencies.py,sha256=KJuAp6Dmv8DQuFnGjbWiHu7StzZj97eBPDyZ_RfCQRc,4141
 datashare_python/discovery.py,sha256=BPB_Ak6d1-vcf9vAQA63IRb2U8h83_mIIi8MbKbFzQ0,7020
 datashare_python/exceptions.py,sha256=bVHEAXxDPKfxeeMC0hJXEsrJkgsKO2ESAhxWU96GA4M,496
-datashare_python/logging_.py,sha256=bhOI5ynVrTe2bU3jUNbaWXfe4zQEG_lLQxJ3Esrz9LQ,2589
+datashare_python/interceptors.py,sha256=Pl7GodPO4KbfflmacpW-vOUgLazjlXSlDNENbpOUt1c,6725
+datashare_python/logging_.py,sha256=-qHz4ztKz4mOCO2z4wunQ4M3xoVhztNvxaiozuMLFRM,2815
 datashare_python/objects.py,sha256=pE0DGNNkl1etxz5ed7T-EaGo1o9TONjH2Lg9u1qdAWU,7571
 datashare_python/task_client.py,sha256=oTmP8bvZW0UyhLNMi1AV3XIAx7hrdbxNRss2Mw2azEc,8435
 datashare_python/template.py,sha256=RxKTYLXoS_EQ8Jc41JkBXppPdbCFqDWfP3BmC0gvB5o,4024
 datashare_python/types_.py,sha256=9Hk1XqpdXbM1TnEzwvJ5G9ABbaCZW9KgBTtiPBVn_7k,649
-datashare_python/utils.py,sha256=ZGZKO9Q4_aLVVilZUCkmHQ21M_37hVOCr7G-qZPOflU,17234
-datashare_python/worker-template.tar.gz,sha256=Krs3td9sVkwXP9KPfIdTUgYWUNoU0G5YLGnn_cNxeTY,280218
-datashare_python/worker.py,sha256=M_I8AL3KZbKeA7-ObaoMUxz3pHBNslJD0XChOcmyafk,7076
+datashare_python/utils.py,sha256=inVjtlBbgL88mN0UM73SSzW76koTW5MGC0NlyopqRW4,17412
+datashare_python/worker-template.tar.gz,sha256=gNSDvn4Lh8iFpFk6j8nlu7pJoIDLh7SVa6EHcjXrj54,286805
+datashare_python/worker.py,sha256=1FdmwYKWKYUKteTM3RC6kFQHR02q8NUDe91hv68QPEo,7207
 datashare_python/cli/__init__.py,sha256=9BPWtssDgsVfWMsZ1TtZCla0EC_kai4RHttr8oNLYOE,1401
 datashare_python/cli/project.py,sha256=w32Gy9AOL5B00uDT4in7YUCt2g68FnNbvwg2M3a8G6o,946
 datashare_python/cli/task.py,sha256=8mvKGS21bZ14BgZ0Uo-dfameljkaI2ZBha80ywCy-E8,5822
 datashare_python/cli/utils.py,sha256=p69CQb0zfixuyBkiZprhdMCc_NuYwXyAn6vC9H1UzAw,911
 datashare_python/cli/worker.py,sha256=I4KTpFIpXFowioFn72Rm6LBCYlY-Dhp4NBIPvtRgUXE,5283
-datashare_python-0.7.0.dist-info/METADATA,sha256=GRZEpX-eTktdvsnSyrvyKzG9kgjqyW4tFs_FDVaEQWU,923
-datashare_python-0.7.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-datashare_python-0.7.0.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
-datashare_python-0.7.0.dist-info/RECORD,,
+datashare_python-0.7.1.dist-info/METADATA,sha256=FF5E62UhxGi0zBseyb4b_6TQRl2JtDZlA1wLECoTblg,923
+datashare_python-0.7.1.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+datashare_python-0.7.1.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
+datashare_python-0.7.1.dist-info/RECORD,,