Adding psycopg2 integration

ext/opentelemetry-ext-postgresql/README.rst

@@ -0,0 +1,25 @@
+OpenTelemetry PostgreSQL integration
+=================================
+
+The integration with PostgreSQL supports the `Psycopg`_ library and is specified
+to ``trace_integration`` using ``'PostgreSQL'``.
+
+.. Psycopg: http://initd.org/psycopg/
+
+Usage
+-----
+
+.. code:: python
+    import psycopg2
+    from opentelemetry.trace import tracer
+    from opentelemetry.trace.ext.postgresql import trace_integration
+    trace_integration(tracer())
+    cnx = psycopg2.connect(database='Database')
+    cursor = cnx.cursor()
+    cursor.execute("INSERT INTO test (testField) VALUES (123)"
+    cursor.close()
+    cnx.close()
+
+References
+----------
+* `OpenTelemetry Project <https://opentelemetry.io/>`_

ext/opentelemetry-ext-postgresql/setup.cfg

@@ -0,0 +1,47 @@
+# Copyright 2019, 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.
+#
+[metadata]
+name = opentelemetry-ext-postgresql
+description = OpenTelemetry PostgreSQL integration
+long_description = file: README.rst
+long_description_content_type = text/x-rst
+author = OpenTelemetry Authors
+author_email = cncf-opentelemetry-contributors@lists.cncf.io
+url = https://github.com/open-telemetry/opentelemetry-python/ext/opentelemetry-ext-postgresql
+platforms = any
+license = Apache-2.0
+classifiers =
+    Development Status :: 3 - Alpha
+    Intended Audience :: Developers
+    License :: OSI Approved :: Apache Software License
+    Programming Language :: Python
+    Programming Language :: Python :: 3
+    Programming Language :: Python :: 3.4
+    Programming Language :: Python :: 3.5
+    Programming Language :: Python :: 3.6
+    Programming Language :: Python :: 3.7
+
+[options]
+python_requires = >=3.4
+package_dir=
+    =src
+packages=find_namespace:
+install_requires =
+    opentelemetry-api >= 0.3.dev0
+    psycopg2-binary >= 2.7.3.1
+    wrapt >= 1.0.0, < 2.0.0
+
+[options.packages.find]
+where = src

ext/opentelemetry-ext-postgresql/setup.py

@@ -0,0 +1,26 @@
+# Copyright 2019, 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 os
+
+import setuptools
+
+BASE_DIR = os.path.dirname(__file__)
+VERSION_FILENAME = os.path.join(
+    BASE_DIR, "src", "opentelemetry", "ext", "postgresql", "version.py"
+)
+PACKAGE_INFO = {}
+with open(VERSION_FILENAME) as f:
+    exec(f.read(), PACKAGE_INFO)
+
+setuptools.setup(version=PACKAGE_INFO["__version__"])

ext/opentelemetry-ext-postgresql/src/opentelemetry/ext/postgresql/__init__.py

@@ -0,0 +1,94 @@
+# Copyright 2019, 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.
+
+"""
+The opentelemetry-ext-postgresql package allows tracing PostgreSQL queries made by the
+Psycopg2 library.
+"""
+
+import typing
+
+import psycopg2
+from psycopg2.extensions import cursor as pgcursor
+import wrapt
+
+from opentelemetry.context import Context
+from opentelemetry.trace import SpanKind, Tracer
+from opentelemetry.trace.status import Status, StatusCanonicalCode
+
+QUERY_WRAP_METHODS = ["execute", "executemany", "callproc"]
+DATABASE_COMPONENT = "postgresql"
+DATABASE_TYPE = "sql"
+
+
+def trace_integration(tracer):
+    """Integrate with PostgreSQL Psycopg library.
+       Psycopg: http://initd.org/psycopg/
+    """
+
+    if tracer is None:
+        raise ValueError("The tracer is not provided.")
+
+    # pylint: disable=unused-argument
+    def wrap(wrapped, instance, args, kwargs):
+        """Patch Psycopg connect method.
+        """
+        connection = wrapped(*args, **kwargs)
+        connection.cursor_factory = TraceCursor
+        return connection
+
+    wrapt.wrap_function_wrapper(psycopg2, "connect", wrap)
+
+    class TraceCursor(pgcursor):
+        def __init__(self, *args, **kwargs):
+            for func in QUERY_WRAP_METHODS:
+                wrapt.wrap_function_wrapper(self, func, self.add_span)
+            super(TraceCursor, self).__init__(*args, **kwargs)
+
+        def add_span(
+            self,
+            wrapped: typing.Callable[..., any],
+            instance: typing.Any,
+            args: typing.Tuple[any, any],
+            kwargs: typing.Dict[any, any],
+        ):
+
+            name = DATABASE_COMPONENT
+            database = self.connection.info.dbname
+            if database:
+                name += "." + database
+
+            query = args[0] if args else ""
+            # Query with parameters
+            if len(args) > 1:
+                query += " params=" + str(args[1])
+
+            with tracer.start_current_span(name, kind=SpanKind.CLIENT) as span:
+                span.set_attribute("component", DATABASE_COMPONENT)
+                span.set_attribute("db.type", DATABASE_TYPE)
+                span.set_attribute("db.instance", database)
+                span.set_attribute("db.statement", query)
+                span.set_attribute("db.user", self.connection.info.user)
+                span.set_attribute("peer.hostname", self.connection.info.host)
+                port = self.connection.info.port
+                if port is not None:
+                    span.set_attribute("peer.port", port)
+                try:
+                    result = wrapped(*args, **kwargs)
+                    span.set_status(Status(StatusCanonicalCode.OK))
+                    return result
+                except Exception as ex:  # pylint: disable=broad-except
+                    span.set_status(
+                        Status(StatusCanonicalCode.UNKNOWN, str(ex))
+                    )

ext/opentelemetry-ext-postgresql/src/opentelemetry/ext/postgresql/version.py

@@ -0,0 +1,15 @@
+# Copyright 2019, 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.
+
+__version__ = "0.3.dev0"

ext/opentelemetry-ext-postgresql/tests/test_postgresql_integration.py

@@ -0,0 +1,31 @@
+# Copyright 2019, 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 unittest
+from unittest import mock
+
+from opentelemetry import trace as trace_api
+
+
+class TestPostgresqlIntegration(unittest.TestCase):
+    def setUp(self):
+        self.tracer = trace_api.tracer()
+        self.start_current_span_patcher = mock.patch.object(
+            self.tracer, "start_as_current_span", autospec=True, spec_set=True
+        )
+
+        self.start_as_current_span = self.start_current_span_patcher.start()
+
+    def tearDown(self):
+        self.start_current_span_patcher.stop()

tox.ini

@@ -2,10 +2,10 @@
 skipsdist = True
 skip_missing_interpreters = True
 envlist =
-    py3{4,5,6,7,8}-test-{api,sdk,example-app,ext-wsgi,ext-flask,ext-http-requests,ext-jaeger,ext-pymongo,opentracing-shim}
-    pypy3-test-{api,sdk,example-app,ext-wsgi,ext-flask,ext-http-requests,ext-jaeger,ext-pymongo,opentracing-shim}
-    py3{4,5,6,7,8}-test-{api,sdk,example-app,example-basic-tracer,example-http,ext-wsgi,ext-flask,ext-http-requests,ext-jaeger,ext-pymongo,opentracing-shim}
-    pypy3-test-{api,sdk,example-app,example-basic-tracer,example-http,ext-wsgi,ext-flask,ext-http-requests,ext-jaeger,ext-pymongo,opentracing-shim}
+    py3{4,5,6,7,8}-test-{api,sdk,example-app,ext-wsgi,ext-flask,ext-http-requests,ext-jaeger,ext-pymongo,ext-postgresql,opentracing-shim}
+    pypy3-test-{api,sdk,example-app,ext-wsgi,ext-flask,ext-http-requests,ext-jaeger,ext-pymongo,ext-postgresql,opentracing-shim}
+    py3{4,5,6,7,8}-test-{api,sdk,example-app,example-basic-tracer,example-http,ext-wsgi,ext-flask,ext-http-requests,ext-jaeger,ext-pymongo,ext-postgresql,opentracing-shim}
+    pypy3-test-{api,sdk,example-app,example-basic-tracer,example-http,ext-wsgi,ext-flask,ext-http-requests,ext-jaeger,ext-pymongo,ext-postgresql,opentracing-shim}
     py3{4,5,6,7,8}-coverage
 
     ; Coverage is temporarily disabled for pypy3 due to the pytest bug.
@@ -36,6 +36,7 @@ changedir =
   test-ext-http-requests: ext/opentelemetry-ext-http-requests/tests
   test-ext-jaeger: ext/opentelemetry-ext-jaeger/tests
   test-ext-pymongo: ext/opentelemetry-ext-pymongo/tests
+  test-ext-postgresql: ext/opentelemetry-ext-postgresql/tests
   test-ext-wsgi: ext/opentelemetry-ext-wsgi/tests
   test-ext-flask: ext/opentelemetry-ext-flask/tests
   test-example-app: examples/opentelemetry-example-app/tests
@@ -66,6 +67,7 @@ commands_pre =
   wsgi,flask: pip install {toxinidir}/ext/opentelemetry-ext-wsgi
   flask: pip install {toxinidir}/ext/opentelemetry-ext-flask[test]
   pymongo: pip install {toxinidir}/ext/opentelemetry-ext-pymongo
+  postgresql: pip install {toxinidir}/ext/opentelemetry-ext-postgresql
   http-requests: pip install {toxinidir}/ext/opentelemetry-ext-http-requests
   jaeger: pip install {toxinidir}/opentelemetry-sdk
   jaeger: pip install {toxinidir}/ext/opentelemetry-ext-jaeger
@@ -115,6 +117,7 @@ commands_pre =
   pip install -e {toxinidir}/ext/opentelemetry-ext-http-requests
   pip install -e {toxinidir}/ext/opentelemetry-ext-jaeger
   pip install -e {toxinidir}/ext/opentelemetry-ext-pymongo
+  pip install -e {toxinidir}/ext/opentelemetry-ext-postgresql
   pip install -e {toxinidir}/ext/opentelemetry-ext-testutil
   pip install -e {toxinidir}/ext/opentelemetry-ext-wsgi
   pip install -e {toxinidir}/ext/opentelemetry-ext-flask[test]
@@ -137,6 +140,8 @@ commands =
          ext/opentelemetry-ext-opentracing-shim/tests/ \
          ext/opentelemetry-ext-pymongo/src/opentelemetry  \
          ext/opentelemetry-ext-pymongo/tests/ \
+         ext/opentelemetry-ext-postgresql/src/opentelemetry  \
+         ext/opentelemetry-ext-postgresql/tests/ \
          ext/opentelemetry-ext-testutil/src/opentelemetry \
          ext/opentelemetry-ext-wsgi/src/ \
          ext/opentelemetry-ext-wsgi/tests/ \