Porting sqlalchemy instrumentation from contrib repo (#591)

Porting the existing sqlalchemy instrumentation from
https://github.com/open-telemetry/opentelemetry-python-contrib/tree/master/reference/ddtrace/contrib/sqlalchemy

Co-authored-by: Mauricio Vásquez <mauricio@kinvolk.io>
Co-authored-by: Diego Hurtado <ocelotl@users.noreply.github.com>
Co-authored-by: Chris Kleinknecht <libc@google.com>

docs/ext/sqlalchemy/sqlalchemy.rst

@@ -0,0 +1,7 @@
+OpenTelemetry SQLAlchemy Instrumentation
+========================================
+
+.. automodule:: opentelemetry.ext.sqlalchemy
+    :members:
+    :undoc-members:
+    :show-inheritance:

ext/opentelemetry-ext-docker-tests/tests/sqlalchemy_tests/__init__.py

@@ -0,0 +1,13 @@
+# 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.

ext/opentelemetry-ext-docker-tests/tests/sqlalchemy_tests/mixins.py

@@ -0,0 +1,184 @@
+# 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 contextlib
+
+from sqlalchemy import Column, Integer, String, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+from opentelemetry import trace
+from opentelemetry.ext.sqlalchemy import SQLAlchemyInstrumentor
+from opentelemetry.ext.sqlalchemy.engine import _DB, _ROWS, _STMT
+from opentelemetry.test.test_base import TestBase
+
+Base = declarative_base()
+
+
+def _create_engine(engine_args):
+    # create a SQLAlchemy engine
+    config = dict(engine_args)
+    url = config.pop("url")
+    return create_engine(url, **config)
+
+
+class Player(Base):
+    """Player entity used to test SQLAlchemy ORM"""
+
+    __tablename__ = "players"
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String(20))
+
+
+class SQLAlchemyTestMixin(TestBase):
+    __test__ = False
+
+    """SQLAlchemy test mixin that includes a complete set of tests
+    that must be executed for different engine. When a new test (or
+    a regression test) should be added to SQLAlchemy test suite, a new
+    entry must be appended here so that it will be executed for all
+    available and supported engines. If the test is specific to only
+    one engine, that test must be added to the specific `TestCase`
+    implementation.
+
+    To support a new engine, create a new `TestCase` that inherits from
+    `SQLAlchemyTestMixin` and `TestCase`. Then you must define the following
+    static class variables:
+    * VENDOR: the database vendor name
+    * SQL_DB: the `db.type` tag that we expect (it's the name of the database available in the `.env` file)
+    * SERVICE: the service that we expect by default
+    * ENGINE_ARGS: all arguments required to create the engine
+
+    To check specific tags in each test, you must implement the
+    `check_meta(self, span)` method.
+    """
+
+    VENDOR = None
+    SQL_DB = None
+    SERVICE = None
+    ENGINE_ARGS = None
+
+    @contextlib.contextmanager
+    def connection(self):
+        # context manager that provides a connection
+        # to the underlying database
+        try:
+            conn = self.engine.connect()
+            yield conn
+        finally:
+            conn.close()
+
+    def check_meta(self, span):
+        """function that can be implemented according to the
+        specific engine implementation
+        """
+
+    def setUp(self):
+        super().setUp()
+        # create an engine with the given arguments
+        self.engine = _create_engine(self.ENGINE_ARGS)
+
+        # create the database / entities and prepare a session for the test
+        Base.metadata.drop_all(bind=self.engine)
+        Base.metadata.create_all(self.engine, checkfirst=False)
+        self.session = sessionmaker(bind=self.engine)()
+        # trace the engine
+        SQLAlchemyInstrumentor().instrument(
+            engine=self.engine, tracer_provider=self.tracer_provider
+        )
+        self.memory_exporter.clear()
+
+    def tearDown(self):
+        # pylint: disable=invalid-name
+        # clear the database and dispose the engine
+        self.session.close()
+        Base.metadata.drop_all(bind=self.engine)
+        self.engine.dispose()
+        SQLAlchemyInstrumentor().uninstrument()
+        super().tearDown()
+
+    def _check_span(self, span):
+        self.assertEqual(span.name, "{}.query".format(self.VENDOR))
+        self.assertEqual(span.attributes.get("service"), self.SERVICE)
+        self.assertEqual(span.attributes.get(_DB), self.SQL_DB)
+        self.assertIs(
+            span.status.canonical_code, trace.status.StatusCanonicalCode.OK
+        )
+        self.assertGreater((span.end_time - span.start_time), 0)
+
+    def test_orm_insert(self):
+        # ensures that the ORM session is traced
+        wayne = Player(id=1, name="wayne")
+        self.session.add(wayne)
+        self.session.commit()
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self._check_span(span)
+        self.assertIn("INSERT INTO players", span.attributes.get(_STMT))
+        self.assertEqual(span.attributes.get(_ROWS), 1)
+        self.check_meta(span)
+
+    def test_session_query(self):
+        # ensures that the Session queries are traced
+        out = list(self.session.query(Player).filter_by(name="wayne"))
+        self.assertEqual(len(out), 0)
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self._check_span(span)
+        self.assertIn(
+            "SELECT players.id AS players_id, players.name AS players_name \nFROM players \nWHERE players.name",
+            span.attributes.get(_STMT),
+        )
+        self.check_meta(span)
+
+    def test_engine_connect_execute(self):
+        # ensures that engine.connect() is properly traced
+        with self.connection() as conn:
+            rows = conn.execute("SELECT * FROM players").fetchall()
+            self.assertEqual(len(rows), 0)
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self._check_span(span)
+        self.assertEqual(span.attributes.get(_STMT), "SELECT * FROM players")
+        self.check_meta(span)
+
+    def test_parent(self):
+        """Ensure that sqlalchemy works with opentelemetry."""
+        tracer = self.tracer_provider.get_tracer("sqlalch_svc")
+
+        with tracer.start_as_current_span("sqlalch_op"):
+            with self.connection() as conn:
+                rows = conn.execute("SELECT * FROM players").fetchall()
+                self.assertEqual(len(rows), 0)
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 2)
+        child_span, parent_span = spans
+
+        # confirm the parenting
+        self.assertIsNone(parent_span.parent)
+        self.assertIs(child_span.parent, parent_span.get_context())
+
+        self.assertEqual(parent_span.name, "sqlalch_op")
+        self.assertEqual(parent_span.instrumentation_info.name, "sqlalch_svc")
+
+        self.assertEqual(child_span.name, "{}.query".format(self.VENDOR))
+        self.assertEqual(child_span.attributes.get("service"), self.SERVICE)

ext/opentelemetry-ext-docker-tests/tests/sqlalchemy_tests/test_instrument.py

@@ -0,0 +1,72 @@
+# 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 os
+import unittest
+
+import sqlalchemy
+
+from opentelemetry import trace
+from opentelemetry.ext.sqlalchemy import SQLAlchemyInstrumentor
+from opentelemetry.test.test_base import TestBase
+
+POSTGRES_CONFIG = {
+    "host": "127.0.0.1",
+    "port": int(os.getenv("TEST_POSTGRES_PORT", "5432")),
+    "user": os.getenv("TEST_POSTGRES_USER", "testuser"),
+    "password": os.getenv("TEST_POSTGRES_PASSWORD", "testpassword"),
+    "dbname": os.getenv("TEST_POSTGRES_DB", "opentelemetry-tests"),
+}
+
+
+class SQLAlchemyInstrumentTestCase(TestBase):
+    """TestCase that checks if the engine is properly traced
+    when the `instrument()` method is used.
+    """
+
+    def setUp(self):
+        # create a traced engine with the given arguments
+        SQLAlchemyInstrumentor().instrument()
+        dsn = (
+            "postgresql://%(user)s:%(password)s@%(host)s:%(port)s/%(dbname)s"
+            % POSTGRES_CONFIG
+        )
+        self.engine = sqlalchemy.create_engine(dsn)
+
+        # prepare a connection
+        self.conn = self.engine.connect()
+        super().setUp()
+
+    def tearDown(self):
+        # clear the database and dispose the engine
+        self.conn.close()
+        self.engine.dispose()
+        SQLAlchemyInstrumentor().uninstrument()
+
+    def test_engine_traced(self):
+        # ensures that the engine is traced
+        rows = self.conn.execute("SELECT 1").fetchall()
+        self.assertEqual(len(rows), 1)
+
+        traces = self.memory_exporter.get_finished_spans()
+        # trace composition
+        self.assertEqual(len(traces), 1)
+        span = traces[0]
+        # check subset of span fields
+        self.assertEqual(span.name, "postgres.query")
+        self.assertEqual(span.attributes.get("service"), "postgres")
+        self.assertIs(
+            span.status.canonical_code, trace.status.StatusCanonicalCode.OK
+        )
+        self.assertGreater((span.end_time - span.start_time), 0)

ext/opentelemetry-ext-docker-tests/tests/sqlalchemy_tests/test_mysql.py

@@ -0,0 +1,77 @@
+# 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 os
+import unittest
+
+import pytest
+from sqlalchemy.exc import ProgrammingError
+
+from opentelemetry import trace
+from opentelemetry.ext.sqlalchemy.engine import _DB, _HOST, _PORT, _ROWS, _STMT
+
+from .mixins import SQLAlchemyTestMixin
+
+MYSQL_CONFIG = {
+    "host": "127.0.0.1",
+    "port": int(os.getenv("TEST_MYSQL_PORT", "3306")),
+    "user": os.getenv("TEST_MYSQL_USER", "testuser"),
+    "password": os.getenv("TEST_MYSQL_PASSWORD", "testpassword"),
+    "database": os.getenv("TEST_MYSQL_DATABASE", "opentelemetry-tests"),
+}
+
+
+class MysqlConnectorTestCase(SQLAlchemyTestMixin):
+    """TestCase for mysql-connector engine"""
+
+    __test__ = True
+
+    VENDOR = "mysql"
+    SQL_DB = "opentelemetry-tests"
+    SERVICE = "mysql"
+    ENGINE_ARGS = {
+        "url": "mysql+mysqlconnector://%(user)s:%(password)s@%(host)s:%(port)s/%(database)s"
+        % MYSQL_CONFIG
+    }
+
+    def check_meta(self, span):
+        # check database connection tags
+        self.assertEqual(span.attributes.get(_HOST), MYSQL_CONFIG["host"])
+        self.assertEqual(span.attributes.get(_PORT), MYSQL_CONFIG["port"])
+
+    def test_engine_execute_errors(self):
+        # ensures that SQL errors are reported
+        with pytest.raises(ProgrammingError):
+            with self.connection() as conn:
+                conn.execute("SELECT * FROM a_wrong_table").fetchall()
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        # span fields
+        self.assertEqual(span.name, "{}.query".format(self.VENDOR))
+        self.assertEqual(span.attributes.get("service"), self.SERVICE)
+        self.assertEqual(
+            span.attributes.get(_STMT), "SELECT * FROM a_wrong_table"
+        )
+        self.assertEqual(span.attributes.get(_DB), self.SQL_DB)
+        self.assertIsNone(span.attributes.get(_ROWS))
+        self.check_meta(span)
+        self.assertTrue(span.end_time - span.start_time > 0)
+        # check the error
+        self.assertIs(
+            span.status.canonical_code,
+            trace.status.StatusCanonicalCode.UNKNOWN,
+        )
+        self.assertIn("a_wrong_table", span.status.description)