Skip to content

Commit

Permalink
Rm instrument_connection is-instrumented check; use dict to store cnx
Browse files Browse the repository at this point in the history
  • Loading branch information
tammy-baylis-swi committed Dec 21, 2024
1 parent a2fa981 commit e8d47bf
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 44 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ class Psycopg2Instrumentor(BaseInstrumentor):

_DATABASE_SYSTEM = "postgresql"

_otel_cursor_factory = None
_otel_cursor_factories = {}

def instrumentation_dependencies(self) -> Collection[str]:
return _instruments
Expand Down Expand Up @@ -193,30 +193,25 @@ def instrument_connection(
Returns:
An instrumented psycopg2 connection object.
"""
if self._is_instrumented_by_opentelemetry:
# _instrument (via BaseInstrumentor) or instrument_connection (this)
# was already called
_logger.warning(
"Attempting to instrument Psycopg2 connection while already instrumented"
)
self._is_instrumented_by_opentelemetry = True
return connection
# TODO Add check for attempt to instrument a connection when already instrumented
# https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3138

# Save cursor_factory at instrumentor level because
# Save cursor_factory in instrumentor map because
# psycopg2 connection type does not allow arbitrary attrs
self._otel_cursor_factory = connection.cursor_factory
self._otel_cursor_factories[connection] = connection.cursor_factory
connection.cursor_factory = _new_cursor_factory(
base_factory=connection.cursor_factory,
tracer_provider=tracer_provider,
)
self._is_instrumented_by_opentelemetry = True
_logger.warning("factories: %s", self._otel_cursor_factories)

return connection

# TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
def uninstrument_connection(self, connection):
self._is_instrumented_by_opentelemetry = False
connection.cursor_factory = self._otel_cursor_factory
connection.cursor_factory = self._otel_cursor_factories.pop(
connection, None
)
return connection


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
from unittest import mock

import psycopg2
import pytest

import opentelemetry.instrumentation.psycopg2
from opentelemetry import trace
Expand Down Expand Up @@ -69,10 +68,6 @@ def get_dsn_parameters(self): # pylint: disable=no-self-use


class TestPostgresqlIntegration(TestBase):
@pytest.fixture(autouse=True)
def inject_fixtures(self, caplog):
self.caplog = caplog # pylint: disable=attribute-defined-outside-init

def setUp(self):
super().setUp()
self.cursor_mock = mock.patch(
Expand Down Expand Up @@ -200,9 +195,8 @@ def test_instrument_connection(self):
instrumentor = Psycopg2Instrumentor()
original_cursor_factory = cnx.cursor_factory
cnx = instrumentor.instrument_connection(cnx)
self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
self.assertEqual(
instrumentor._otel_cursor_factory, original_cursor_factory
instrumentor._otel_cursor_factories[cnx], original_cursor_factory
)
cursor = cnx.cursor()
cursor.execute(query)
Expand All @@ -222,22 +216,27 @@ def test_instrument_connection_with_instrument(self):

instrumentor = Psycopg2Instrumentor()
instrumentor.instrument()
self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)

cnx = psycopg2.connect(database="test")
original_cursor_factory = cnx.cursor_factory
orig_cnx = cnx
cnx = Psycopg2Instrumentor().instrument_connection(cnx)

# TODO Should not set cursor_factory if already instrumented
# https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3138
# self.assertEqual(instrumentor._otel_cursor_factories[cnx], None)
self.assertEqual(
self.caplog.records[0].getMessage(),
"Attempting to instrument Psycopg2 connection while already instrumented",
instrumentor._otel_cursor_factories.get(orig_cnx),
original_cursor_factory,
)
self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
# Will not set cursor_factory if already instrumented
self.assertEqual(instrumentor._otel_cursor_factory, None)
cursor = cnx.cursor()
cursor.execute(query)

spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
# TODO Add check for attempt to instrument a connection when already instrumented
# https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3138
# self.assertEqual(len(spans_list), 1)
self.assertEqual(len(spans_list), 2)

def test_instrument_connection_with_instrument_connection(self):
cnx = psycopg2.connect(database="test")
Expand All @@ -251,26 +250,29 @@ def test_instrument_connection_with_instrument_connection(self):
cnx = psycopg2.connect(database="test")
original_cursor_factory = cnx.cursor_factory
instrumentor = Psycopg2Instrumentor()
orig_cnx = cnx
cnx = instrumentor.instrument_connection(cnx)
self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
self.assertEqual(
instrumentor._otel_cursor_factory, original_cursor_factory
instrumentor._otel_cursor_factories.get(orig_cnx),
original_cursor_factory,
)

cnx = Psycopg2Instrumentor().instrument_connection(cnx)
self.assertEqual(
self.caplog.records[0].getMessage(),
"Attempting to instrument Psycopg2 connection while already instrumented",
)
self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
orig_cnx = cnx
original_cursor_factory = cnx.cursor_factory
instrumentor = Psycopg2Instrumentor()
cnx = instrumentor.instrument_connection(cnx)
self.assertEqual(
instrumentor._otel_cursor_factory, original_cursor_factory
instrumentor._otel_cursor_factories.get(orig_cnx),
original_cursor_factory,
)
cursor = cnx.cursor()
cursor.execute(query)

spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
# TODO Add check for attempt to instrument a connection when already instrumented
# https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3138
# self.assertEqual(len(spans_list), 1)
self.assertEqual(len(spans_list), 2)

# pylint: disable=unused-argument
def test_uninstrument_connection_with_instrument(self):
Expand All @@ -283,15 +285,13 @@ def test_uninstrument_connection_with_instrument(self):

spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)

cnx = Psycopg2Instrumentor().uninstrument_connection(cnx)
cursor = cnx.cursor()
cursor.execute(query)

spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
self.assertFalse(instrumentor._is_instrumented_by_opentelemetry)

# pylint: disable=unused-argument
def test_uninstrument_connection_with_instrument_connection(self):
Expand All @@ -305,19 +305,20 @@ def test_uninstrument_connection_with_instrument_connection(self):

spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
self.assertEqual(
instrumentor._otel_cursor_factory, original_cursor_factory
instrumentor._otel_cursor_factories[cnx], original_cursor_factory
)

orig_cnx = cnx
cnx = Psycopg2Instrumentor().uninstrument_connection(cnx)
cursor = cnx.cursor()
cursor.execute(query)

spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
self.assertFalse(instrumentor._is_instrumented_by_opentelemetry)
self.assertEqual(instrumentor._otel_cursor_factory, None)
self.assertEqual(
instrumentor._otel_cursor_factories.get(orig_cnx), None
)

@mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
def test_sqlcommenter_enabled(self, event_mocked):
Expand Down

0 comments on commit e8d47bf

Please sign in to comment.