Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix psycopg2 instrument issue #2795

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
([#2823](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2823))
- `opentelemetry-instrumentation-system-metrics` fix `process.runtime.cpu.utilization` values to be shown in range of 0 to 1
([#2812](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2812))
- `opentelemetry-instrumentation-psycopg2` Fix psycopg2 instrument issue
([#2795](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2795))
- `opentelemetry-instrumentation-fastapi` fix `fastapi` auto-instrumentation by removing `fastapi-slim` support, `fastapi-slim` itself is discontinued from maintainers
([#2783](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2783))
- `opentelemetry-instrumentation-aws-lambda` Avoid exception when a handler is not present.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@
cursor as pg_cursor, # pylint: disable=no-name-in-module
)
from psycopg2.sql import Composed # pylint: disable=no-name-in-module
from wrapt import ObjectProxy

from opentelemetry.instrumentation import dbapi
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
Expand Down Expand Up @@ -158,24 +159,25 @@ def _uninstrument(self, **kwargs):
dbapi.unwrap_connect(psycopg2, "connect")

# TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
@staticmethod
def instrument_connection(connection, tracer_provider=None):
if not hasattr(connection, "_is_instrumented_by_opentelemetry"):
lzchen marked this conversation as resolved.
Show resolved Hide resolved
connection._is_instrumented_by_opentelemetry = False

if not connection._is_instrumented_by_opentelemetry:
setattr(
connection, _OTEL_CURSOR_FACTORY_KEY, connection.cursor_factory
)
connection.cursor_factory = _new_cursor_factory(
tracer_provider=tracer_provider
)
connection._is_instrumented_by_opentelemetry = True
else:
def instrument_connection(self, connection, tracer_provider=None):
if isinstance(connection, ObjectProxy):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm still not sure this solves the original issue (we still don't know WHY the error is being generated when we are trying to access _is_instrumented_by_opentelemetry field but this current change is fine as a short-term solution. Please add a TODO or follow-up issue to investigate.

_logger.warning(
"Attempting to instrument Psycopg connection while already instrumented"
)
return connection
return connection

db_integration = DatabaseApiIntegration(
__name__,
self._DATABASE_SYSTEM,
connection_attributes=self._CONNECTION_ATTRIBUTES,
version=__version__,
tracer_provider=tracer_provider,
)
db_integration.get_connection_attributes(connection)
db_integration.cursor_factory = _new_cursor_factory(
tracer_provider=tracer_provider
)
return dbapi.get_traced_connection_proxy(connection, db_integration)

# TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
@staticmethod
Expand All @@ -184,7 +186,7 @@ def uninstrument_connection(connection):
connection, _OTEL_CURSOR_FACTORY_KEY, None
)

return connection
return dbapi.uninstrument_connection(connection)


# TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,7 @@ def test_instrument_connection(self):
self.assertEqual(len(spans_list), 0)

cnx = Psycopg2Instrumentor().instrument_connection(cnx)
self.assertFalse(hasattr(cnx, "_is_instrumented_by_opentelemetry"))
cursor = cnx.cursor()
cursor.execute(query)

Expand All @@ -209,6 +210,7 @@ def test_instrument_connection_with_instrument(self):

Psycopg2Instrumentor().instrument()
cnx = Psycopg2Instrumentor().instrument_connection(cnx)
self.assertFalse(hasattr(cnx, "_is_instrumented_by_opentelemetry"))
cursor = cnx.cursor()
cursor.execute(query)

Expand Down Expand Up @@ -236,7 +238,7 @@ def test_uninstrument_connection_with_instrument(self):
# pylint: disable=unused-argument
def test_uninstrument_connection_with_instrument_connection(self):
cnx = psycopg2.connect(database="test")
Psycopg2Instrumentor().instrument_connection(cnx)
cnx = Psycopg2Instrumentor().instrument_connection(cnx)
query = "SELECT * FROM test"
cursor = cnx.cursor()
cursor.execute(query)
Expand Down Expand Up @@ -281,3 +283,30 @@ def test_no_op_tracer_provider(self):
cursor.execute(query)
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 0)

# pylint: disable=unused-argument
def test_duplicated_instrumentation_works(self):
cnx = psycopg2.connect(database="test")
query = "SELECT * FROM test"
cursor = cnx.cursor()
cursor.execute(query)

spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 0)

Psycopg2Instrumentor().instrument()
Psycopg2Instrumentor().instrument()
cnx = Psycopg2Instrumentor().instrument_connection(cnx)
self.assertFalse(hasattr(cnx, "_is_instrumented_by_opentelemetry"))
cursor = cnx.cursor()
cursor.execute(query)

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

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