From e861b93362126f5c701982efcf43c5d2e99f2f50 Mon Sep 17 00:00:00 2001 From: Thiyagu55 <64461612+Thiyagu55@users.noreply.github.com> Date: Fri, 18 Mar 2022 06:49:30 +0530 Subject: [PATCH] Extending SQLCommenter support of dbapi into psycopg2 (#940) * Added configuration for sqlcommenter * Fixed linting errors * Updated CHANGELOG.md * Unit test case for sqlcommenter in psycopg2 * Merge conflict resolved * Update CHANGELOG.md * psycopg2 sqlcommenter integration test * linting changes * linting changes * linting changes * linting changes * linting changes * version compatability issue fix Co-authored-by: Srikanth Chekuri Co-authored-by: Leighton Chen --- CHANGELOG.md | 5 +- .../instrumentation/psycopg2/__init__.py | 3 +- .../tests/test_psycopg2_integration.py | 20 +++++++ .../tests/postgres/test_psycopg_functional.py | 25 ++++++++ .../postgres/test_psycopg_sqlcommenter.py | 59 +++++++++++++++++++ 5 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b70b522960..3a3851461c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.10.0-0.29b0...HEAD) +### Added + +- `opentelemetry-instrumentation-psycopg2` extended the sql commenter support of dbapi into psycopg2 + ([#940](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/940)) - `opentelemetry-instrumentation-flask` Fix non-recording span bug ([#999])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/999) - `opentelemetry-instrumentation-tornado` Fix non-recording span bug @@ -32,7 +36,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ([#908](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/908)) - `opentelemetry-instrumentation-requests` make span attribute available to samplers ([#931](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/931)) - - `opentelemetry-datadog-exporter` add deprecation note to example. ([#900](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/900)) diff --git a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py index e338902e1b..4b7ff23ea6 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py @@ -76,7 +76,7 @@ def _instrument(self, **kwargs): Psycopg: http://initd.org/psycopg/ """ tracer_provider = kwargs.get("tracer_provider") - + enable_sqlcommenter = kwargs.get("enable_commenter", False) dbapi.wrap_connect( __name__, psycopg2, @@ -86,6 +86,7 @@ def _instrument(self, **kwargs): version=__version__, tracer_provider=tracer_provider, db_api_integration_factory=DatabaseApiIntegration, + enable_commenter=enable_sqlcommenter, ) def _uninstrument(self, **kwargs): diff --git a/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py index d20cc41348..e974da7d60 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py +++ b/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py @@ -224,3 +224,23 @@ def test_uninstrument_connection_with_instrument_connection(self): spans_list = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans_list), 1) + + @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect") + def test_sqlcommenter_enabled(self, event_mocked): + cnx = psycopg2.connect(database="test") + Psycopg2Instrumentor().instrument(enable_commenter=True) + query = "SELECT * FROM test" + cursor = cnx.cursor() + cursor.execute(query) + kwargs = event_mocked.call_args[1] + self.assertEqual(kwargs["enable_commenter"], True) + + @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect") + def test_sqlcommenter_disabled(self, event_mocked): + cnx = psycopg2.connect(database="test") + Psycopg2Instrumentor().instrument() + query = "SELECT * FROM test" + cursor = cnx.cursor() + cursor.execute(query) + kwargs = event_mocked.call_args[1] + self.assertEqual(kwargs["enable_commenter"], False) diff --git a/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_functional.py b/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_functional.py index e007bb7e84..38ecff02e7 100644 --- a/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_functional.py +++ b/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_functional.py @@ -155,3 +155,28 @@ def test_composed_queries(self): span.attributes[SpanAttributes.DB_STATEMENT], 'SELECT FROM "users" where "name"=\'"abc"\'', ) + + def test_commenter_enabled(self): + + stmt = "CREATE TABLE IF NOT EXISTS users (id integer, name varchar)" + with self._tracer.start_as_current_span("rootSpan"): + self._cursor.execute(stmt) + self.validate_spans("CREATE") + Psycopg2Instrumentor().uninstrument() + Psycopg2Instrumentor().instrument(enable_commenter=True) + + self._cursor.execute( + sql.SQL("SELECT FROM {table} where {field}='{value}'").format( + table=sql.Identifier("users"), + field=sql.Identifier("name"), + value=sql.Identifier("abc"), + ) + ) + + spans = self.memory_exporter.get_finished_spans() + span = spans[2] + self.assertEqual(span.name, "SELECT") + self.assertEqual( + span.attributes[SpanAttributes.DB_STATEMENT], + 'SELECT FROM "users" where "name"=\'"abc"\'', + ) diff --git a/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py b/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py new file mode 100644 index 0000000000..1e3460600e --- /dev/null +++ b/tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py @@ -0,0 +1,59 @@ +# Copyright 2020, 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 psycopg2 +from test_psycopg_functional import ( + POSTGRES_DB_NAME, + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, +) + +from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor +from opentelemetry.test.test_base import TestBase + + +class TestFunctionalPsycopg(TestBase): + @classmethod + def setUpClass(cls): + super().setUpClass() + cls._connection = None + cls._cursor = None + cls._tracer = cls.tracer_provider.get_tracer(__name__) + Psycopg2Instrumentor().instrument(enable_commenter=True) + cls._connection = psycopg2.connect( + dbname=POSTGRES_DB_NAME, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + host=POSTGRES_HOST, + port=POSTGRES_PORT, + ) + cls._connection.set_session(autocommit=True) + cls._cursor = cls._connection.cursor() + + @classmethod + def tearDownClass(cls): + if cls._cursor: + cls._cursor.close() + if cls._connection: + cls._connection.close() + Psycopg2Instrumentor().uninstrument() + + def test_commenter_enabled(self): + self._cursor.execute("SELECT 1;") + self.assertRegex( + self._cursor.query.decode("ascii"), + r"SELECT 1; /\*traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/", + )