Fix pyodbc cursor error in SQLA instrumentation (#469)

CHANGELOG.md

@@ -6,11 +6,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.3.0-0.22b0...HEAD)
 
+### Changed
 - `opentelemetry-instrumentation-tornado` properly instrument work done in tornado on_finish method.
   ([#499](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/499))
 - `opentelemetry-instrumentation` Fixed cases where trying to use an instrumentation package without the
   target library was crashing auto instrumentation agent.
   ([#530](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/530))
+- Fix weak reference error for pyodbc cursor in SQLAlchemy instrumentation.
+  ([#469](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/469))
 
 ## [0.22b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.3.0-0.22b0) - 2021-06-01
 

instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 from threading import local
-from weakref import WeakKeyDictionary
 
 from sqlalchemy.event import listen  # pylint: disable=no-name-in-module
 
@@ -60,7 +59,7 @@ class EngineTracer:
         self.tracer = tracer
         self.engine = engine
         self.vendor = _normalize_vendor(engine.name)
-        self.cursor_mapping = WeakKeyDictionary()
+        self.cursor_mapping = {}
         self.local = local()
 
         listen(engine, "before_cursor_execute", self._before_cur_exec)
@@ -116,6 +115,7 @@ class EngineTracer:
             return
 
         span.end()
+        self._cleanup(cursor)
 
     def _handle_error(self, context):
         span = self.current_thread_span
@@ -129,6 +129,13 @@ class EngineTracer:
                 )
         finally:
             span.end()
+            self._cleanup(context.cursor)
+
+    def _cleanup(self, cursor):
+        try:
+            del self.cursor_mapping[cursor]
+        except KeyError:
+            pass
 
 
 def _get_attributes_from_url(url):

tests/opentelemetry-docker-tests/tests/check_availability.py

@@ -18,6 +18,7 @@ import time
 import mysql.connector
 import psycopg2
 import pymongo
+import pyodbc
 import redis
 
 MONGODB_COLLECTION_NAME = "test"
@@ -36,6 +37,11 @@ POSTGRES_PORT = int(os.getenv("POSTGRESQL_PORT", "5432"))
 POSTGRES_USER = os.getenv("POSTGRESQL_USER", "testuser")
 REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
 REDIS_PORT = int(os.getenv("REDIS_PORT ", "6379"))
+MSSQL_DB_NAME = os.getenv("MSSQL_DB_NAME", "opentelemetry-tests")
+MSSQL_HOST = os.getenv("MSSQL_HOST", "localhost")
+MSSQL_PORT = int(os.getenv("MSSQL_PORT", "1433"))
+MSSQL_USER = os.getenv("MSSQL_USER", "sa")
+MSSQL_PASSWORD = os.getenv("MSSQL_PASSWORD", "yourStrong(!)Password")
 RETRY_COUNT = 8
 RETRY_INTERVAL = 5  # Seconds
 
@@ -104,12 +110,23 @@ def check_redis_connection():
     connection.hgetall("*")
 
 
+@retryable
+def check_mssql_connection():
+    connection = pyodbc.connect(
+        f"DRIVER={{ODBC Driver 17 for SQL Server}};SERVER={MSSQL_HOST},"
+        f"{MSSQL_PORT};DATABASE={MSSQL_DB_NAME};UID={MSSQL_USER};"
+        f"PWD={MSSQL_PASSWORD}"
+    )
+    connection.close()
+
+
 def check_docker_services_availability():
     # Check if Docker services accept connections
     check_pymongo_connection()
     check_mysql_connection()
     check_postgres_connection()
     check_redis_connection()
+    check_mssql_connection()
 
 
 check_docker_services_availability()

tests/opentelemetry-docker-tests/tests/docker-compose.yml

@@ -39,3 +39,11 @@ services:
         - "16686:16686"
         - "14268:14268"
         - "9411:9411"
+  otmssql:
+    image: mcr.microsoft.com/mssql/server:2017-latest
+    ports:
+      - "1433:1433"
+    environment:
+      ACCEPT_EULA: "Y"
+      SA_PASSWORD: "yourStrong(!)Password"
+    command: /bin/sh -c "sleep 10s && /opt/mssql-tools/bin/sqlcmd -S localhost -U sa -P yourStrong\(!\)Password -d master -Q 'CREATE DATABASE [opentelemetry-tests]' & /opt/mssql/bin/sqlservr"

tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mssql.py

@@ -0,0 +1,107 @@
+# 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.semconv.trace import SpanAttributes
+
+from .mixins import Player, SQLAlchemyTestMixin
+
+MSSQL_CONFIG = {
+    "host": "127.0.0.1",
+    "port": int(os.getenv("TEST_MSSQL_PORT", "1433")),
+    "user": os.getenv("TEST_MSSQL_USER", "sa"),
+    "password": os.getenv("TEST_MSSQL_PASSWORD", "yourStrong(!)Password"),
+    "database": os.getenv("TEST_MSSQL_DATABASE", "opentelemetry-tests"),
+    "driver": os.getenv("TEST_MSSQL_DRIVER", "ODBC+Driver+17+for+SQL+Server"),
+}
+
+
+class MssqlConnectorTestCase(SQLAlchemyTestMixin):
+    """TestCase for pyodbc engine"""
+
+    __test__ = True
+
+    VENDOR = "mssql"
+    SQL_DB = "opentelemetry-tests"
+    ENGINE_ARGS = {
+        "url": "mssql+pyodbc://%(user)s:%(password)s@%(host)s:%(port)s/%(database)s?driver=%(driver)s"
+        % MSSQL_CONFIG
+    }
+
+    def check_meta(self, span):
+        # check database connection tags
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.NET_PEER_NAME),
+            MSSQL_CONFIG["host"],
+        )
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.NET_PEER_PORT),
+            MSSQL_CONFIG["port"],
+        )
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.DB_NAME),
+            MSSQL_CONFIG["database"],
+        )
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.DB_USER), MSSQL_CONFIG["user"]
+        )
+
+    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, "SELECT opentelemetry-tests")
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.DB_STATEMENT),
+            "SELECT * FROM a_wrong_table",
+        )
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.DB_NAME), self.SQL_DB
+        )
+        self.check_meta(span)
+        self.assertTrue(span.end_time - span.start_time > 0)
+        # check the error
+        self.assertIs(
+            span.status.status_code, trace.StatusCode.ERROR,
+        )
+        self.assertIn("a_wrong_table", span.status.description)
+
+    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()
+        # identity insert on before the insert, insert, and identity insert off after the insert
+        self.assertEqual(len(spans), 3)
+        span = spans[1]
+        self._check_span(span, "INSERT")
+        self.assertIn(
+            "INSERT INTO players",
+            span.attributes.get(SpanAttributes.DB_STATEMENT),
+        )
+        self.check_meta(span)

tox.ini

@@ -404,6 +404,7 @@ deps =
   celery[pytest] >= 4.0, < 6.0
   protobuf>=3.13.0
   requests==2.25.0
+  pyodbc~=4.0.30
 changedir =
   tests/opentelemetry-docker-tests/tests