Merge branch 'main' into pymemcache_instruments_version_bump

srikanthccv · web-flow · commit 48df656ff22a · 2022-03-09T18:15:23.000+05:30
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,8 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.9.1-0.28b1...HEAD)
 
+- `opentelemetry-instrumentation-wsgi` Capture custom request/response headers in span attributes
+  ([#925])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/925)
+
+- `opentelemetry-instrumentation-flask` Flask: Capture custom request/response headers in span attributes
+  ([#952])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/952)
+
 ### Added
 
+- `opentelemetry-instrumentation-sqlalchemy` added experimental sql commenter capability
+   ([#924](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/924))
 - `opentelemetry-instrumentation-dbapi` add experimental sql commenter capability
   ([#908](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/908))
 - `opentelemetry-instrumentation-requests` make span attribute available to samplers
diff --git a/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py b/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py
@@ -275,7 +275,7 @@ def process_response(self, request, response):
                 add_response_attributes(
                     span,
                     f"{response.status_code} {response.reason_phrase}",
-                    response,
+                    response.items(),
                 )
 
             propagator = get_global_response_propagator()
diff --git a/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py b/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py
@@ -153,6 +153,10 @@ def _start_response(status, response_headers, *args, **kwargs):
                     otel_wsgi.add_response_attributes(
                         span, status, response_headers
                     )
+                    if span.kind == trace.SpanKind.SERVER:
+                        otel_wsgi.add_custom_response_headers(
+                            span, response_headers
+                        )
                 else:
                     _logger.warning(
                         "Flask environ's OpenTelemetry span "
@@ -200,6 +204,10 @@ def _before_request():
                 ] = flask.request.url_rule.rule
             for key, value in attributes.items():
                 span.set_attribute(key, value)
+            if span.kind == trace.SpanKind.SERVER:
+                otel_wsgi.add_custom_request_headers(
+                    span, flask_request_environ
+                )
 
         activation = trace.use_span(span, end_on_exit=True)
         activation.__enter__()  # pylint: disable=E1101
diff --git a/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py b/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from flask import Response
 from werkzeug.test import Client
 from werkzeug.wrappers import BaseResponse
 
@@ -23,6 +24,16 @@ def _hello_endpoint(helloid):
             raise ValueError(":-(")
         return "Hello: " + str(helloid)
 
+    @staticmethod
+    def _custom_response_headers():
+        resp = Response("test response")
+        resp.headers["content-type"] = "text/plain; charset=utf-8"
+        resp.headers["content-length"] = "13"
+        resp.headers[
+            "my-custom-header"
+        ] = "my-custom-value-1,my-custom-header-2"
+        return resp
+
     def _common_initialization(self):
         def excluded_endpoint():
             return "excluded"
@@ -35,6 +46,9 @@ def excluded2_endpoint():
         self.app.route("/excluded/<int:helloid>")(self._hello_endpoint)
         self.app.route("/excluded")(excluded_endpoint)
         self.app.route("/excluded2")(excluded2_endpoint)
+        self.app.route("/test_custom_response_headers")(
+            self._custom_response_headers
+        )
 
         # pylint: disable=attribute-defined-outside-init
         self.client = Client(self.app, BaseResponse)
diff --git a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py
@@ -442,3 +442,101 @@ def test_mark_span_internal_in_presence_of_span_from_other_framework(self):
         self.assertEqual(
             span_list[0].parent.span_id, span_list[1].context.span_id
         )
+
+
+class TestCustomRequestResponseHeaders(
+    InstrumentationTest, TestBase, WsgiTestBase
+):
+    def setUp(self):
+        super().setUp()
+
+        self.env_patch = patch.dict(
+            "os.environ",
+            {
+                "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST": "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3",
+                "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE": "content-type,content-length,my-custom-header,invalid-header",
+            },
+        )
+        self.env_patch.start()
+        self.app = Flask(__name__)
+        FlaskInstrumentor().instrument_app(self.app)
+
+        self._common_initialization()
+
+    def tearDown(self):
+        super().tearDown()
+        self.env_patch.stop()
+        with self.disable_logging():
+            FlaskInstrumentor().uninstrument_app(self.app)
+
+    def test_custom_request_header_added_in_server_span(self):
+        headers = {
+            "Custom-Test-Header-1": "Test Value 1",
+            "Custom-Test-Header-2": "TestValue2,TestValue3",
+        }
+        resp = self.client.get("/hello/123", headers=headers)
+        self.assertEqual(200, resp.status_code)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.request.header.custom_test_header_1": ("Test Value 1",),
+            "http.request.header.custom_test_header_2": (
+                "TestValue2,TestValue3",
+            ),
+        }
+        self.assertEqual(span.kind, trace.SpanKind.SERVER)
+        self.assertSpanHasAttributes(span, expected)
+
+    def test_custom_request_header_not_added_in_internal_span(self):
+        tracer = trace.get_tracer(__name__)
+        with tracer.start_as_current_span("test", kind=trace.SpanKind.SERVER):
+            headers = {
+                "Custom-Test-Header-1": "Test Value 1",
+                "Custom-Test-Header-2": "TestValue2,TestValue3",
+            }
+            resp = self.client.get("/hello/123", headers=headers)
+            self.assertEqual(200, resp.status_code)
+            span = self.memory_exporter.get_finished_spans()[0]
+            not_expected = {
+                "http.request.header.custom_test_header_1": ("Test Value 1",),
+                "http.request.header.custom_test_header_2": (
+                    "TestValue2,TestValue3",
+                ),
+            }
+            self.assertEqual(span.kind, trace.SpanKind.INTERNAL)
+            for key, _ in not_expected.items():
+                self.assertNotIn(key, span.attributes)
+
+    def test_custom_response_header_added_in_server_span(self):
+        resp = self.client.get("/test_custom_response_headers")
+        self.assertEqual(resp.status_code, 200)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.response.header.content_type": (
+                "text/plain; charset=utf-8",
+            ),
+            "http.response.header.content_length": ("13",),
+            "http.response.header.my_custom_header": (
+                "my-custom-value-1,my-custom-header-2",
+            ),
+        }
+        self.assertEqual(span.kind, trace.SpanKind.SERVER)
+        self.assertSpanHasAttributes(span, expected)
+
+    def test_custom_response_header_not_added_in_internal_span(self):
+        tracer = trace.get_tracer(__name__)
+        with tracer.start_as_current_span("test", kind=trace.SpanKind.SERVER):
+            resp = self.client.get("/test_custom_response_headers")
+            self.assertEqual(resp.status_code, 200)
+            span = self.memory_exporter.get_finished_spans()[0]
+            not_expected = {
+                "http.response.header.content_type": (
+                    "text/plain; charset=utf-8",
+                ),
+                "http.response.header.content_length": ("13",),
+                "http.response.header.my_custom_header": (
+                    "my-custom-value-1,my-custom-header-2",
+                ),
+            }
+            self.assertEqual(span.kind, trace.SpanKind.INTERNAL)
+            for key, _ in not_expected.items():
+                self.assertNotIn(key, span.attributes)
diff --git a/instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py b/instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py
@@ -161,7 +161,7 @@ def trace_tween(request):
                 otel_wsgi.add_response_attributes(
                     span,
                     response_or_exception.status,
-                    response_or_exception.headers,
+                    response_or_exception.headerlist,
                 )
 
                 propagator = get_global_response_propagator()
diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py
@@ -106,6 +106,7 @@ def _instrument(self, **kwargs):
             return EngineTracer(
                 _get_tracer(kwargs.get("engine"), tracer_provider),
                 kwargs.get("engine"),
+                kwargs.get("enable_commenter", False),
             )
         return None
 
diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
@@ -17,7 +17,12 @@
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.sqlalchemy.version import __version__
+from opentelemetry.instrumentation.utils import (
+    _generate_opentelemetry_traceparent,
+    _generate_sql_comment,
+)
 from opentelemetry.semconv.trace import NetTransportValues, SpanAttributes
+from opentelemetry.trace import Span
 from opentelemetry.trace.status import Status, StatusCode
 
 
@@ -70,12 +75,15 @@ def _wrap_create_engine_internal(func, module, args, kwargs):
 
 
 class EngineTracer:
-    def __init__(self, tracer, engine):
+    def __init__(self, tracer, engine, enable_commenter=False):
         self.tracer = tracer
         self.engine = engine
         self.vendor = _normalize_vendor(engine.name)
+        self.enable_commenter = enable_commenter
 
-        listen(engine, "before_cursor_execute", self._before_cur_exec)
+        listen(
+            engine, "before_cursor_execute", self._before_cur_exec, retval=True
+        )
         listen(engine, "after_cursor_execute", _after_cur_exec)
         listen(engine, "handle_error", _handle_error)
 
@@ -115,6 +123,18 @@ def _before_cur_exec(
                     span.set_attribute(key, value)
 
         context._otel_span = span
+        if self.enable_commenter:
+            statement = statement + EngineTracer._generate_comment(span=span)
+
+        return statement, params
+
+    @staticmethod
+    def _generate_comment(span: Span) -> str:
+        span_context = span.get_span_context()
+        meta = {}
+        if span_context.is_valid:
+            meta.update(_generate_opentelemetry_traceparent(span))
+        return _generate_sql_comment(**meta)
 
 
 # pylint: disable=unused-argument
diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import asyncio
+import logging
 from unittest import mock
 
 import pytest
@@ -20,12 +21,17 @@
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+from opentelemetry.instrumentation.sqlalchemy.engine import EngineTracer
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider, export
 from opentelemetry.test.test_base import TestBase
 
 
 class TestSqlalchemyInstrumentation(TestBase):
+    @pytest.fixture(autouse=True)
+    def inject_fixtures(self, caplog):
+        self.caplog = caplog  # pylint: disable=attribute-defined-outside-init
+
     def tearDown(self):
         super().tearDown()
         SQLAlchemyInstrumentor().uninstrument()
@@ -150,3 +156,22 @@ async def run():
             self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
 
         asyncio.get_event_loop().run_until_complete(run())
+
+    def test_generate_commenter(self):
+        logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+            enable_commenter=True,
+        )
+
+        cnx = engine.connect()
+        cnx.execute("SELECT	1 + 1;").fetchall()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertIn(
+            EngineTracer._generate_comment(span),
+            self.caplog.records[-2].getMessage(),
+        )
diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlcommenter.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlcommenter.py
@@ -0,0 +1,53 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from sqlalchemy import create_engine
+
+from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+from opentelemetry.test.test_base import TestBase
+
+
+class TestSqlalchemyInstrumentationWithSQLCommenter(TestBase):
+    @pytest.fixture(autouse=True)
+    def inject_fixtures(self, caplog):
+        self.caplog = caplog  # pylint: disable=attribute-defined-outside-init
+
+    def tearDown(self):
+        super().tearDown()
+        SQLAlchemyInstrumentor().uninstrument()
+
+    def test_sqlcommenter_enabled(self):
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+            enable_commenter=True,
+        )
+        cnx = engine.connect()
+        cnx.execute("SELECT  1;").fetchall()
+        self.assertRegex(
+            self.caplog.records[-2].getMessage(),
+            r"SELECT  1; /\*traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/",
+        )
+
+    def test_sqlcommenter_disabled(self):
+        engine = create_engine("sqlite:///:memory:", echo=True)
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine, tracer_provider=self.tracer_provider
+        )
+        cnx = engine.connect()
+        cnx.execute("SELECT 1;").fetchall()
+
+        self.assertEqual(self.caplog.records[-2].getMessage(), "SELECT 1;")
diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py b/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py
diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/utils.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/utils.py
diff --git a/util/opentelemetry-util-http/src/opentelemetry/util/http/__init__.py b/util/opentelemetry-util-http/src/opentelemetry/util/http/__init__.py
diff --git a/util/opentelemetry-util-http/tests/test_capture_custom_headers.py b/util/opentelemetry-util-http/tests/test_capture_custom_headers.py

Original file line number	Diff line number	Diff line change
`@@ -275,7 +275,7 @@ def process_response(self, request, response):`
`275`	`275`	`add_response_attributes(`
`276`	`276`	`span,`
`277`	`277`	`f"{response.status_code} {response.reason_phrase}",`
`278`		`- response,`
	`278`	`+ response.items(),`
`279`	`279`	`)`
`280`	`280`
`281`	`281`	`propagator = get_global_response_propagator()`
Original file line number	Diff line number	Diff line change
`@@ -161,7 +161,7 @@ def trace_tween(request):`
`161`	`161`	`otel_wsgi.add_response_attributes(`
`162`	`162`	`span,`
`163`	`163`	`response_or_exception.status,`
`164`		`- response_or_exception.headers,`
	`164`	`+ response_or_exception.headerlist,`
`165`	`165`	`)`
`166`	`166`
`167`	`167`	`propagator = get_global_response_propagator()`
Original file line number	Diff line number	Diff line change
`@@ -106,6 +106,7 @@ def _instrument(self, **kwargs):`
`106`	`106`	`return EngineTracer(`
`107`	`107`	`_get_tracer(kwargs.get("engine"), tracer_provider),`
`108`	`108`	`kwargs.get("engine"),`
	`109`	`+ kwargs.get("enable_commenter", False),`
`109`	`110`	`)`
`110`	`111`	`return None`
`111`	`112`