datashare-python 0.6.3__tar.gz → 0.7.1__tar.gz

{datashare_python-0.6.3 → datashare_python-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.6.3
+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/
@@ -8,13 +8,14 @@ Project-URL: Repository, https://github.com/ICIJ/datashare-python
 Project-URL: Issues, https://github.com/ICIJ/datashare-python/issues
 Author-email: Clément Doumouro <cdoumouro@icij.org>, Clément Doumouro <clement.doumouro@gmail.com>, Lion Summerbell <lsummerbell@icij.org>
 Requires-Python: <4,>=3.11
-Requires-Dist: aiohttp~=3.11.9
-Requires-Dist: alive-progress~=3.2.0
-Requires-Dist: hatchling~=1.27.0
+Requires-Dist: aiohttp~=3.11
+Requires-Dist: alive-progress~=3.2
+Requires-Dist: hatchling~=1.27
 Requires-Dist: icij-common[elasticsearch]~=0.8.2
-Requires-Dist: nest-asyncio~=1.6.0
-Requires-Dist: python-json-logger~=4.0.0
+Requires-Dist: nest-asyncio~=1.6
+Requires-Dist: orjson~=3.11
+Requires-Dist: python-json-logger~=4.0
 Requires-Dist: pyyaml~=6.0
-Requires-Dist: temporalio~=1.23.0
+Requires-Dist: temporalio~=1.23
 Requires-Dist: tomlkit~=0.14.0
 Requires-Dist: typer~=0.15.4

{datashare_python-0.6.3 → datashare_python-0.7.1}/datashare_python/config.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import ClassVar
+from typing import Annotated, Literal
 
 from icij_common.es import ESClient
 from icij_common.pydantic_utils import ICIJSettings
@@ -18,7 +18,6 @@ import datashare_python
 from .objects import BaseModel
 from .task_client import DatashareTaskClient
 from .types_ import TemporalClient
-from .utils import LogWithWorkerIDMixin
 
 _ALL_LOGGERS = [datashare_python.__name__]
 
@@ -76,11 +75,20 @@ class TemporalClientConfig(BaseModel):
         return self._client
 
 
-class WorkerConfig(ICIJSettings, LogWithWorkerIDMixin, BaseModel):
+LogLevel = Literal["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"]
+
+
+class LoggingConfig(BaseModel):
+    log_in_json: bool = False
+    loggers: dict[str, LogLevel]
+
+
+class WorkerConfig(ICIJSettings, BaseModel):
     model_config = DS_WORKER_SETTINGS_CONFIG
 
-    loggers: ClassVar[list[str]] = Field(_ALL_LOGGERS, frozen=True)
-    log_level: str = Field(default="INFO")
+    logging: Annotated[LoggingConfig, Field(frozen=True)] = {
+        datashare_python.__name__: "INFO"
+    }
 
     datashare: DatashareClientConfig = DatashareClientConfig()
     elasticsearch: ESClientConfig = ESClientConfig()

{datashare_python-0.6.3 → datashare_python-0.7.1}/datashare_python/conftest.py

@@ -8,10 +8,12 @@ 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 (
     DatashareClientConfig,
+    LoggingConfig,
     TemporalClientConfig,
     WorkerConfig,
 )
@@ -19,10 +21,7 @@ from datashare_python.dependencies import (
     lifespan_es_client,
     lifespan_task_client,
     set_es_client,
-    set_event_loop,
-    set_loggers,
     set_task_client,
-    set_temporal_client,
     with_dependencies,
 )
 from datashare_python.objects import Document, TaskState
@@ -78,13 +77,7 @@ class MockedWorkflow:
 
 @pytest.fixture(scope="session")
 def test_deps() -> list[ContextManagerFactory]:
-    return [
-        set_loggers,
-        set_event_loop,
-        set_es_client,
-        set_temporal_client,
-        set_task_client,
-    ]
+    return [set_es_client, set_task_client]
 
 
 @pytest.fixture(scope="session")
@@ -99,8 +92,16 @@ def event_loop(
 
 @pytest.fixture(scope="session")
 def test_worker_config() -> WorkerConfig:
+    logging_config = LoggingConfig(
+        log_in_json=False,
+        loggers={
+            "datashare_python": "DEBUG",
+            "icij_common": "DEBUG",
+            "worker_template": "DEBUG",
+        },
+    )
     return WorkerConfig(
-        log_level="DEBUG",
+        logging=logging_config,
         datashare=DatashareClientConfig(url="http://localhost:8080"),
         temporal=TemporalClientConfig(host="localhost:7233"),
     )

{datashare_python-0.6.3 → datashare_python-0.7.1}/datashare_python/dependencies.py

@@ -9,8 +9,9 @@ from typing import Any
 
 from icij_common.es import ESClient
 
-from .config import WorkerConfig
+from .config import LogLevel, WorkerConfig
 from .exceptions import DependencyInjectionError
+from .logging_ import setup_worker_loggers
 from .task_client import DatashareTaskClient
 from .types_ import ContextManagerFactory, TemporalClient
 
@@ -35,10 +36,13 @@ def lifespan_event_loop() -> AbstractEventLoop:
         raise DependencyInjectionError("event loop") from e
 
 
-def set_loggers(worker_config: WorkerConfig, worker_id: str) -> None:
-    worker_config.setup_loggers(worker_id=worker_id)
+def set_loggers(
+    worker_config: WorkerConfig, worker_id: str, loggers: dict[str, LogLevel]
+) -> None:
+    setup_worker_loggers(
+        loggers=loggers, worker_id=worker_id, in_json=worker_config.logging.log_in_json
+    )
     logger.info("worker loggers ready to log 💬")
-    logger.info("app config: %s", worker_config.model_dump_json(indent=2))
 
 
 def set_worker_config(worker_config: WorkerConfig) -> None:

{datashare_python-0.6.3 → datashare_python-0.7.1}/datashare_python/discovery.py

@@ -4,7 +4,7 @@ from collections.abc import Callable, Iterable
 from importlib.metadata import entry_points
 
 from .config import WorkerConfig
-from .dependencies import set_worker_config
+from .dependencies import set_loggers, set_worker_config
 from .types_ import ContextManagerFactory
 from .utils import ActivityWithProgress
 
@@ -29,6 +29,8 @@ _Discovery = tuple[
     type[WorkerConfig],
 ]
 
+_MANDATORY_DEPS = [set_worker_config, set_loggers]
+
 
 def discover(
     wf_names: list[str] | None, *, act_names: list[str] | None, deps_name: str | None
@@ -68,8 +70,9 @@ def discover(
     deps = []
     if deps_name is not None:
         deps = discover_dependencies(deps_name)
-    if set_worker_config not in deps:
-        deps.append(set_worker_config)
+    for mandatory in _MANDATORY_DEPS:
+        if mandatory not in deps:
+            deps.append(mandatory)
     if deps:
         n_deps = len(deps)
         discovered += "\n"

datashare_python-0.7.1/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-0.7.1/datashare_python/logging_.py

@@ -0,0 +1,86 @@
+import logging
+import sys
+from copy import copy
+
+from icij_common.logging_utils import (
+    DATE_FMT,
+    STREAM_HANDLER_FMT,
+    STREAM_HANDLER_FMT_WITH_WORKER_ID,
+)
+from pythonjsonlogger.core import RESERVED_ATTRS, BaseJsonFormatter
+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"]
+_TRACE_CONTEXT_ATTRS = ["trace_id", "parent_id", "traceparent"]
+_LOGGED_ATTRIBUTES = (
+    copy(RESERVED_ATTRS)
+    + _WF_LOGGED_ATTRS
+    + _ACT_LOGGER_ATTRS
+    + _TRACE_CONTEXT_ATTRS
+    + ["worker_id"]
+)
+
+
+def setup_worker_loggers(
+    loggers: dict[str, LogLevel], *, worker_id: str | None, in_json: bool
+) -> None:
+    worker_filter = WorkerFilter(worker_id)
+    for logger_name, level_str in loggers.items():
+        level = getattr(logging, level_str)
+        logger = logging.getLogger(logger_name)
+        logger.setLevel(level)
+        logger.handlers = []
+        for handler in _get_worker_handlers(level, worker_id, in_json=in_json):
+            logger.addHandler(handler)
+        logger.addFilter(worker_filter)
+
+
+def _get_worker_handlers(
+    level: int, worker_id: str | None, *, in_json: bool
+) -> list[logging.Handler]:
+    stream_handler = logging.StreamHandler(sys.stderr)
+    if in_json:
+        fmt = _json_formatter(datefmt=DATE_FMT)
+    else:
+        if worker_id is not None:
+            fmt = STREAM_HANDLER_FMT_WITH_WORKER_ID
+        else:
+            fmt = STREAM_HANDLER_FMT
+        fmt = logging.Formatter(fmt, DATE_FMT)
+    stream_handler.setFormatter(fmt)
+    stream_handler.setLevel(level)
+    return [stream_handler]
+
+
+class WorkerFilter(logging.Filter):
+    def __init__(self, worker_id: str) -> None:
+        super().__init__()
+        self._worker_id = worker_id
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        record.worker_id = self._worker_id
+        if workflow.in_workflow():
+            wf_info = workflow.info()
+            for attr in _WF_LOGGED_ATTRS:
+                setattr(record, attr, getattr(wf_info, attr))
+        if activity.in_activity():
+            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) -> BaseJsonFormatter:
+    fmt = OrjsonFormatter(  # let's keep logging as fast as possible
+        _LOGGED_ATTRIBUTES, datefmt=datefmt
+    )
+    return fmt

{datashare_python-0.6.3 → datashare_python-0.7.1}/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-0.6.3 → datashare_python-0.7.1}/datashare_python/worker.py

@@ -8,6 +8,7 @@ from asyncio import AbstractEventLoop
 from collections.abc import AsyncGenerator, Callable
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
+from copy import copy
 from typing import Any
 
 from temporalio.worker import PollerBehaviorSimpleMaximum, Worker
@@ -15,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__)
@@ -83,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,
@@ -142,12 +145,27 @@ async def worker_context(
     task_queue: str,
     dependencies: list[ContextManagerFactory] | None = None,
 ) -> AsyncGenerator[DatashareWorker, None]:
+    discovered = []
+    if activities is not None:
+        discovered.extend(activities)
+    if workflows is not None:
+        discovered.extend(workflows)
+    if dependencies is not None:
+        discovered.extend(dependencies)
+    discovered.append(worker_config)
+    loggers = copy(worker_config.logging.loggers)
+    discovered_loggers = {_get_object_package(o).__name__ for o in discovered}
+    for logger in discovered_loggers:
+        if logger not in loggers:
+            # Log in info by default
+            loggers[logger] = "INFO"
     deps_cm = (
         with_dependencies(
             dependencies,
             worker_config=worker_config,
             worker_id=worker_id,
             event_loop=event_loop,
+            loggers=loggers,
         )
         if dependencies
         else _do_nothing_cm()
@@ -181,3 +199,9 @@ def _get_class_from_method(method: Callable) -> type:
     class_name = method.__qualname__.rsplit(".", 1)[0]
     module = sys.modules[method.__module__]
     return getattr(module, class_name)
+
+
+def _get_object_package(obj: Any) -> Any:
+    mod = inspect.getmodule(obj)
+    base, _, _ = mod.__name__.partition(".")
+    return sys.modules[base]

{datashare_python-0.6.3 → datashare_python-0.7.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datashare-python"
-version = "0.6.3"
+version = "0.7.1"
 description = "Manage Pythoœn tasks and local resources in Datashare"
 authors = [
   { name = "Clément Doumouro", email = "cdoumouro@icij.org" },
@@ -10,16 +10,17 @@ authors = [
 readme = "README.md"
 requires-python = ">=3.11,<4"
 dependencies = [
-  "alive-progress~=3.2.0",
-  "aiohttp~=3.11.9",
+  "alive-progress~=3.2",
+  "aiohttp~=3.11",
   "icij-common[elasticsearch]~=0.8.2",
-  "python-json-logger~=4.0.0",
-  "nest-asyncio~=1.6.0",
-  "temporalio~=1.23.0",
+  "python-json-logger~=4.0",
+  "nest-asyncio~=1.6",
+  "temporalio~=1.23",
   "typer~=0.15.4",
   "tomlkit~=0.14.0",
-  "hatchling~=1.27.0",
+  "hatchling~=1.27",
   "pyyaml~=6.0",
+  "orjson~=3.11",
 ]
 
 [project.urls]