Using Resources and overloading the Op for telemerty?

[whillas-yabble]

Hi,
We're trying to figure out how to pass around telemerty spans automatically in each Op (don't think there is much difference here with Assets?). Got some code going what kinda works but requires each @op to take a resource key which is going to be much work for us as there are hundreds of ops across all our pipelines. Here is some code I got working:

import functools
import inspect

import logfire

from dagster import (
    ConfigurableResource,
    Definitions,
    InitResourceContext,
    OpExecutionContext,
    job,
)
from dagster import op as dagster_op
from opentelemetry import context as otel_context
from opentelemetry.trace import (
    NonRecordingSpan,
    SpanContext,
    TraceFlags,
    set_span_in_context,
)

# Configure logfire (Replace with your actual Logfire/OTel setup)

logfire.configure(
    service_name="testing-123", send_to_logfire="if-token-present", environment="local"
)


class DistributedTracingResource(ConfigurableResource):
    """
    Derives a deterministic OpenTelemetry Trace ID from the Dagster Run ID.
    Ensures all ops in a multiprocess run share the same parent trace.
    """

    def setup_for_execution(self, context: InitResourceContext) -> None:
        if not context.run_id:
            return

        # 1. Convert Dagster Run ID (UUID) to an OTel 128-bit Trace ID
        trace_id = int(context.run_id.replace("-", ""), 16)

        # 2. Create a deterministic 64-bit Span ID for the "Parent" (the run itself)
        span_id = int(context.run_id.replace("-", "")[:16], 16)

        # 3. Create the OpenTelemetry SpanContext
        span_context = SpanContext(
            trace_id=trace_id,
            span_id=span_id,
            is_remote=True,  # Indicates this parent context came from "outside"
            trace_flags=TraceFlags(1),  # Mark it as sampled/recorded
        )

        # 4. Inject this parent context into the current Python process
        parent_span = NonRecordingSpan(span_context)
        otel_context.attach(set_span_in_context(parent_span))

        context.log.info(f"Initialized OTel Trace ID matching Run ID: {context.run_id}")


def traced_op(*op_args, **op_kwargs):
    def decorator(func):
        sig = inspect.signature(func)
        user_wants_context = "context" in sig.parameters

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # 1. Isolate the Dagster context
            ctx = kwargs.get("context")
            if not ctx:
                for arg in args:
                    if isinstance(arg, OpExecutionContext):
                        ctx = arg
                        break

            try:
                # 2. Open the Logfire Span (It will automatically inherit the Trace ID from the Resource)
                op_name = ctx.op.name if ctx else func.__name__
                with logfire.span(f"op: {op_name}"):
                    # 3. Clean up args before calling the user's original function
                    if not user_wants_context:
                        kwargs.pop("context", None)
                        args = tuple(
                            a for a in args if not isinstance(a, OpExecutionContext)
                        )

                    return func(*args, **kwargs)
            finally:
                pass  # The Resource's teardown handles OTel context detachment if needed

        # 4. Signature Magic: Force Dagster to inject `context` if the user didn't ask for it
        if not user_wants_context:
            parameters = list(sig.parameters.values())
            context_param = inspect.Parameter(
                "context",
                inspect.Parameter.POSITIONAL_OR_KEYWORD,
                annotation=OpExecutionContext,
            )
            parameters.insert(0, context_param)
            wrapper.__signature__ = sig.replace(parameters=parameters)
        else:
            wrapper.__signature__ = sig

        return dagster_op(*op_args, **op_kwargs)(wrapper)

    # Handle usage without parentheses: @traced_op vs @traced_op(name="foo")
    if len(op_args) == 1 and callable(op_args[0]) and not op_kwargs:
        func = op_args[0]
        op_args = ()
        return decorator(func)

    return decorator


# --- 1. Define Ops using the new decorator ---


@traced_op(required_resource_keys={"tracing_resource"})
def fetch_data():
    # logfire will auto-nest this print statement under the `op: fetch_data` span!
    logfire.info("Fetching data from API...")
    return [1, 2, 3, 4, 5]


@traced_op(
    name="process_records_op", required_resource_keys={"tracing_resource"}
)  # Custom names are supported!
def process_data(context, data: list):
    logfire.info(f"Processing {len(data)} records")
    return [x * 2 for x in data]


@traced_op(required_resource_keys={"tracing_resource"})
def save_data(data: list):
    logfire.info(f"Saving data: {data}")


# --- 2. Build the Job ---


@job(resource_defs={"tracing_resource": DistributedTracingResource()})
def my_traced_pipeline():
    raw_data = fetch_data()
    processed_data = process_data(raw_data)
    save_data(processed_data)


# --- 3. Bind the Resource to the Definitions ---


def main():
    # This is how we'll deploy it
    # defs = Definitions(
    #     jobs=[my_traced_pipeline],
    #     resources={
    #         # Inject the tracing resource. Dagster will run this once per op process.
    #         "tracing_resource": DistributedTracingResource()
    #     },
    # )

    my_traced_pipeline.execute_in_process(
        resources={"tracing_resource": DistributedTracingResource()}
    )  


if __name__ == "__main__":
    main()

The question are two:

1. Is there a better way to do this natively in Dagster?
2. Is there a way to get this custom decorator to use the resource automatically without having to manually supply the required_resource_keys every time?

This gives use nested spans like this: