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 <srikanth.chekuri92@gmail.com>
Co-authored-by: Leighton Chen <lechen@microsoft.com>

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))
 

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):

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)

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"\'',
+        )

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}'\*/",
+        )