Repository: apache/arrow
File: python/pyarrow/tests/test_flight.py (Python)

Code before:
        # trace_context will be evaluated after we've exited the scope
        # of the OTel span (and so the value we want won't be present)
        return ((f"{key}: {value}").encode("utf-8")
                for (key, value) in trace_context.items())


def test_tracing():
    with TracingFlightServer(middleware={
            "tracing": flight.TracingServerMiddlewareFactory(),
    }) as server, \
            FlightClient(('localhost', server.port)) as client:
        # We can't tell if Arrow was built with OpenTelemetry support,
        # so we can't count on any particular values being there; we
        # can only ensure things don't blow up either way.
        options = flight.FlightCallOptions(headers=[
            # Pretend we have an OTel implementation
            (b"traceparent", b"00-000ff00f00f0ff000f0f00ff0f00fff0-"
                             b"000f0000f0f00000-00"),
            (b"tracestate", b""),
        ])
        for value in client.do_action((b"", b""), options=options):
            pass


@pytest.mark.skipif(sys.platform == "win32",
                    reason="Opentelemetry traces don't appear on stdout on Windows")
def test_tracing_server_middleware_emits_traces():
    # Validate that we are able to emit traces to stdout when
    # ARROW_TRACING_BACKEND=ostream
    if not is_opentelemetry_enabled():
        pytest.skip("Arrow not built with OpenTelemetry")
    code = """if 1:
        import pyarrow.flight as flight
        from pyarrow.flight import FlightServerBase, FlightClient

        class SimpleServer(FlightServerBase):
            def do_action(self, context, action):
                return []

        with SimpleServer(
            middleware={"otel": flight.TracingServerMiddlewareFactory()}
        ) as server:
            with FlightClient(('localhost', server.port)) as client:
                list(client.do_action((b"", b"")))
        """
    env = os.environ.copy()
    env['ARROW_TRACING_BACKEND'] = "ostream"
    env['OTEL_SERVICE_NAME'] = "pyarrow-testing-service"
    res = subprocess.run([sys.executable, "-c", code], env=env,
                         capture_output=True)
    assert res.returncode == 0, res.stderr

Review comment (suggestion):
```suggestion
# More info about the skip: GH-49447
@pytest.mark.skipif(sys.platform == "win32",
```

Code after:
        # trace_context will be evaluated after we've exited the scope
        # of the OTel span (and so the value we want won't be present)
        return ((f"{key}: {value}").encode("utf-8")
                for (key, value) in trace_context.items())


def test_tracing():
    with TracingFlightServer(middleware={
            "tracing": flight.TracingServerMiddlewareFactory(),
    }) as server, \
            FlightClient(('localhost', server.port)) as client:
        # We can't tell if Arrow was built with OpenTelemetry support,
        # so we can't count on any particular values being there; we
        # can only ensure things don't blow up either way.
        options = flight.FlightCallOptions(headers=[
            # Pretend we have an OTel implementation
            (b"traceparent", b"00-000ff00f00f0ff000f0f00ff0f00fff0-"
                             b"000f0000f0f00000-00"),
            (b"tracestate", b""),
        ])
        for value in client.do_action((b"", b""), options=options):
            pass


# More info about the skip: GH-49447
@pytest.mark.skipif(sys.platform == "win32",
                    reason="Opentelemetry traces don't appear on stdout on Windows")
def test_tracing_server_middleware_emits_traces():
    # Validate that we are able to emit traces to stdout when
    # ARROW_TRACING_BACKEND=ostream
    if not is_opentelemetry_enabled():
        pytest.skip("Arrow not built with OpenTelemetry")
    code = """if 1:
        import pyarrow.flight as flight
        from pyarrow.flight import FlightServerBase, FlightClient

        class SimpleServer(FlightServerBase):
            def do_action(self, context, action):
                return []

        with SimpleServer(
            middleware={"otel": flight.TracingServerMiddlewareFactory()}
        ) as server:
            with FlightClient(('localhost', server.port)) as client:
                list(client.do_action((b"", b"")))
        """
    env = os.environ.copy()
    env['ARROW_TRACING_BACKEND'] = "ostream"
    env['OTEL_SERVICE_NAME'] = "pyarrow-testing-service"
    res = subprocess.run([sys.executable, "-c", code], env=env,
                         capture_output=True)