Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 2d9c8df

Browse files
cnnradamscnnradams
authored
ext/psycopg2: Implement BaseInstrumentor interface (open-telemetry#694)
- Implemented BaseInstrumentor interface to enable auto-instrumentation - Added integration tests (same tests as other db integrations)
1 parent 9e58b8a commit 2d9c8df

File tree

8 files changed

+216
-128
lines changed

8 files changed

+216
-128
lines changed

‎ext/opentelemetry-ext-dbapi/src/opentelemetry/ext/dbapi/__init__.py

Lines changed: 37 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ def instrument_connection(
174174
connection_attributes=connection_attributes,
175175
)
176176
db_integration.get_connection_attributes(connection)
177-
return TracedConnectionProxy(connection, db_integration)
177+
return get_traced_connection_proxy(connection, db_integration)
178178

179179

180180
def uninstrument_connection(connection):
@@ -227,7 +227,7 @@ def wrapped_connection(
227227
"""
228228
connection = connect_method(*args, **kwargs)
229229
self.get_connection_attributes(connection)
230-
return TracedConnectionProxy(connection, self)
230+
return get_traced_connection_proxy(connection, self)
231231

232232
def get_connection_attributes(self, connection):
233233
# Populate span fields using connection
@@ -260,23 +260,21 @@ def get_connection_attributes(self, connection):
260260
self.span_attributes["net.peer.port"] = port
261261

262262

263-
# pylint: disable=abstract-method
264-
class TracedConnectionProxy(wrapt.ObjectProxy):
265-
# pylint: disable=unused-argument
266-
def __init__(
267-
self,
268-
connection,
269-
db_api_integration: DatabaseApiIntegration,
270-
*args,
271-
**kwargs
272-
):
273-
wrapt.ObjectProxy.__init__(self, connection)
274-
self._db_api_integration = db_api_integration
263+
def get_traced_connection_proxy(
264+
connection, db_api_integration, *args, **kwargs
265+
):
266+
# pylint: disable=abstract-method
267+
class TracedConnectionProxy(wrapt.ObjectProxy):
268+
# pylint: disable=unused-argument
269+
def __init__(self, connection, *args, **kwargs):
270+
wrapt.ObjectProxy.__init__(self, connection)
271+
272+
def cursor(self, *args, **kwargs):
273+
return get_traced_cursor_proxy(
274+
self.__wrapped__.cursor(*args, **kwargs), db_api_integration
275+
)
275276

276-
def cursor(self, *args, **kwargs):
277-
return TracedCursorProxy(
278-
self.__wrapped__.cursor(*args, **kwargs), self._db_api_integration
279-
)
277+
return TracedConnectionProxy(connection, *args, **kwargs)
280278

281279

282280
class TracedCursor:
@@ -323,31 +321,28 @@ def traced_execution(
323321
raise ex
324322

325323

326-
# pylint: disable=abstract-method
327-
class TracedCursorProxy(wrapt.ObjectProxy):
324+
def get_traced_cursor_proxy(cursor, db_api_integration, *args, **kwargs):
325+
_traced_cursor = TracedCursor(db_api_integration)
326+
# pylint: disable=abstract-method
327+
class TracedCursorProxy(wrapt.ObjectProxy):
328328

329-
# pylint: disable=unused-argument
330-
def __init__(
331-
self,
332-
cursor,
333-
db_api_integration: DatabaseApiIntegration,
334-
*args,
335-
**kwargs
336-
):
337-
wrapt.ObjectProxy.__init__(self, cursor)
338-
self._traced_cursor = TracedCursor(db_api_integration)
329+
# pylint: disable=unused-argument
330+
def __init__(self, cursor, *args, **kwargs):
331+
wrapt.ObjectProxy.__init__(self, cursor)
339332

340-
def execute(self, *args, **kwargs):
341-
return self._traced_cursor.traced_execution(
342-
self.__wrapped__.execute, *args, **kwargs
343-
)
333+
def execute(self, *args, **kwargs):
334+
return _traced_cursor.traced_execution(
335+
self.__wrapped__.execute, *args, **kwargs
336+
)
344337

345-
def executemany(self, *args, **kwargs):
346-
return self._traced_cursor.traced_execution(
347-
self.__wrapped__.executemany, *args, **kwargs
348-
)
338+
def executemany(self, *args, **kwargs):
339+
return _traced_cursor.traced_execution(
340+
self.__wrapped__.executemany, *args, **kwargs
341+
)
349342

350-
def callproc(self, *args, **kwargs):
351-
return self._traced_cursor.traced_execution(
352-
self.__wrapped__.callproc, *args, **kwargs
353-
)
343+
def callproc(self, *args, **kwargs):
344+
return _traced_cursor.traced_execution(
345+
self.__wrapped__.callproc, *args, **kwargs
346+
)
347+
348+
return TracedCursorProxy(cursor, *args, **kwargs)

‎ext/opentelemetry-ext-dbapi/tests/test_dbapi_integration.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -125,15 +125,13 @@ def test_wrap_connect(self, mock_dbapi):
125125
dbapi.wrap_connect(self.tracer, mock_dbapi, "connect", "-")
126126
connection = mock_dbapi.connect()
127127
self.assertEqual(mock_dbapi.connect.call_count, 1)
128-
self.assertIsInstance(connection, dbapi.TracedConnectionProxy)
129128
self.assertIsInstance(connection.__wrapped__, mock.Mock)
130129

131130
@mock.patch("opentelemetry.ext.dbapi")
132131
def test_unwrap_connect(self, mock_dbapi):
133132
dbapi.wrap_connect(self.tracer, mock_dbapi, "connect", "-")
134133
connection = mock_dbapi.connect()
135134
self.assertEqual(mock_dbapi.connect.call_count, 1)
136-
self.assertIsInstance(connection, dbapi.TracedConnectionProxy)
137135

138136
dbapi.unwrap_connect(mock_dbapi, "connect")
139137
connection = mock_dbapi.connect()
@@ -145,7 +143,6 @@ def test_instrument_connection(self):
145143
# Avoid get_attributes failing because can't concatenate mock
146144
connection.database = "-"
147145
connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
148-
self.assertIsInstance(connection2, dbapi.TracedConnectionProxy)
149146
self.assertIs(connection2.__wrapped__, connection)
150147

151148
def test_uninstrument_connection(self):
@@ -154,7 +151,6 @@ def test_uninstrument_connection(self):
154151
# be concatenated
155152
connection.database = "-"
156153
connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
157-
self.assertIsInstance(connection2, dbapi.TracedConnectionProxy)
158154
self.assertIs(connection2.__wrapped__, connection)
159155

160156
connection3 = dbapi.uninstrument_connection(connection2)

‎ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
import psycopg2
1919

2020
from opentelemetry import trace as trace_api
21-
from opentelemetry.ext.psycopg2 import trace_integration
21+
from opentelemetry.ext.psycopg2 import Psycopg2Instrumentor
2222
from opentelemetry.test.test_base import TestBase
2323

2424
POSTGRES_HOST = os.getenv("POSTGRESQL_HOST ", "localhost")
@@ -35,7 +35,7 @@ def setUpClass(cls):
3535
cls._connection = None
3636
cls._cursor = None
3737
cls._tracer = cls.tracer_provider.get_tracer(__name__)
38-
trace_integration(cls.tracer_provider)
38+
Psycopg2Instrumentor().instrument(tracer_provider=cls.tracer_provider)
3939
cls._connection = psycopg2.connect(
4040
dbname=POSTGRES_DB_NAME,
4141
user=POSTGRES_USER,
@@ -52,6 +52,7 @@ def tearDownClass(cls):
5252
cls._cursor.close()
5353
if cls._connection:
5454
cls._connection.close()
55+
Psycopg2Instrumentor().uninstrument()
5556

5657
def validate_spans(self):
5758
spans = self.memory_exporter.get_finished_spans()

‎ext/opentelemetry-ext-psycopg2/CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
## Unreleased
44

5+
- Implement instrumentor interface, enabling auto-instrumentation ([#694]https://github.com/open-telemetry/opentelemetry-python/pull/694)
6+
57
## 0.4a0
68

79
Released 2020-02-21

‎ext/opentelemetry-ext-psycopg2/setup.cfg

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,18 @@ package_dir=
4141
packages=find_namespace:
4242
install_requires =
4343
opentelemetry-api == 0.8.dev0
44+
opentelemetry-ext-dbapi == 0.8.dev0
45+
opentelemetry-auto-instrumentation == 0.8.dev0
4446
psycopg2-binary >= 2.7.3.1
4547
wrapt >= 1.0.0, < 2.0.0
4648

49+
[options.extras_require]
50+
test =
51+
opentelemetry-test == 0.8.dev0
52+
4753
[options.packages.find]
4854
where = src
55+
56+
[options.entry_points]
57+
opentelemetry_instrumentor =
58+
psycopg2 = opentelemetry.ext.psycopg2:Psycopg2Instrumentor

‎ext/opentelemetry-ext-psycopg2/src/opentelemetry/ext/psycopg2/__init__.py

Lines changed: 66 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@
1313
# limitations under the License.
1414

1515
"""
16-
The integration with PostgreSQL supports the `Psycopg`_ library and is specified
17-
to ``trace_integration`` using ``'PostgreSQL'``.
16+
The integration with PostgreSQL supports the `Psycopg`_ library, it can be enabled by
17+
using ``Psycopg2Instrumentor``.
1818
1919
.. _Psycopg: http://initd.org/psycopg/
2020
@@ -26,11 +26,12 @@
2626
import psycopg2
2727
from opentelemetry import trace
2828
from opentelemetry.sdk.trace import TracerProvider
29-
from opentelemetry.trace.ext.psycopg2 import trace_integration
29+
from opentelemetry.trace.ext.psycopg2 import Psycopg2Instrumentor
3030
3131
trace.set_tracer_provider(TracerProvider())
3232
33-
trace_integration()
33+
Psycopg2Instrumentor().instrument()
34+
3435
cnx = psycopg2.connect(database='Database')
3536
cursor = cnx.cursor()
3637
cursor.execute("INSERT INTO test (testField) VALUES (123)")
@@ -41,83 +42,77 @@
4142
---
4243
"""
4344

44-
import logging
4545
import typing
4646

4747
import psycopg2
4848
import wrapt
49-
from psycopg2.sql import Composable
5049

51-
from opentelemetry.ext.dbapi import DatabaseApiIntegration, TracedCursor
50+
from opentelemetry.auto_instrumentation.instrumentor import BaseInstrumentor
51+
from opentelemetry.ext import dbapi
5252
from opentelemetry.ext.psycopg2.version import __version__
53-
from opentelemetry.trace import Tracer, get_tracer
54-
55-
logger = logging.getLogger(__name__)
56-
57-
DATABASE_COMPONENT = "postgresql"
58-
DATABASE_TYPE = "sql"
53+
from opentelemetry.trace import TracerProvider, get_tracer
5954

6055

61-
def trace_integration(tracer_provider=None):
62-
"""Integrate with PostgreSQL Psycopg library.
63-
Psycopg: http://initd.org/psycopg/
64-
"""
65-
66-
tracer = get_tracer(__name__, __version__, tracer_provider)
67-
68-
connection_attributes = {
56+
class Psycopg2Instrumentor(BaseInstrumentor):
57+
_CONNECTION_ATTRIBUTES = {
6958
"database": "info.dbname",
7059
"port": "info.port",
7160
"host": "info.host",
7261
"user": "info.user",
7362
}
74-
db_integration = DatabaseApiIntegration(
75-
tracer,
76-
DATABASE_COMPONENT,
77-
database_type=DATABASE_TYPE,
78-
connection_attributes=connection_attributes,
79-
)
80-
81-
# pylint: disable=unused-argument
82-
def wrap_connect(
83-
connect_func: typing.Callable[..., any],
84-
instance: typing.Any,
85-
args: typing.Tuple[any, any],
86-
kwargs: typing.Dict[any, any],
87-
):
88-
connection = connect_func(*args, **kwargs)
89-
db_integration.get_connection_attributes(connection)
90-
connection.cursor_factory = PsycopgTraceCursor
91-
return connection
92-
93-
try:
94-
wrapt.wrap_function_wrapper(psycopg2, "connect", wrap_connect)
95-
except Exception as ex: # pylint: disable=broad-except
96-
logger.warning("Failed to integrate with pyscopg2. %s", str(ex))
97-
98-
class PsycopgTraceCursor(psycopg2.extensions.cursor):
99-
def __init__(self, *args, **kwargs):
100-
self._traced_cursor = TracedCursor(db_integration)
101-
super(PsycopgTraceCursor, self).__init__(*args, **kwargs)
102-
103-
# pylint: disable=redefined-builtin
104-
def execute(self, query, vars=None):
105-
if isinstance(query, Composable):
106-
query = query.as_string(self)
107-
return self._traced_cursor.traced_execution(
108-
super(PsycopgTraceCursor, self).execute, query, vars
109-
)
110-
111-
# pylint: disable=redefined-builtin
112-
def executemany(self, query, vars):
113-
if isinstance(query, Composable):
114-
query = query.as_string(self)
115-
return self._traced_cursor.traced_execution(
116-
super(PsycopgTraceCursor, self).executemany, query, vars
117-
)
118-
119-
# pylint: disable=redefined-builtin
120-
def callproc(self, procname, vars=None):
121-
return self._traced_cursor.traced_execution(
122-
super(PsycopgTraceCursor, self).callproc, procname, vars
123-
)
63+
64+
_DATABASE_COMPONENT = "postgresql"
65+
_DATABASE_TYPE = "sql"
66+
67+
def _instrument(self, **kwargs):
68+
"""Integrate with PostgreSQL Psycopg library.
69+
Psycopg: http://initd.org/psycopg/
70+
"""
71+
72+
tracer_provider = kwargs.get("tracer_provider")
73+
74+
tracer = get_tracer(__name__, __version__, tracer_provider)
75+
76+
dbapi.wrap_connect(
77+
tracer,
78+
psycopg2,
79+
"connect",
80+
self._DATABASE_COMPONENT,
81+
self._DATABASE_TYPE,
82+
self._CONNECTION_ATTRIBUTES,
83+
)
84+
85+
def _uninstrument(self, **kwargs):
86+
""""Disable Psycopg2 instrumentation"""
87+
dbapi.unwrap_connect(psycopg2, "connect")
88+
89+
# pylint:disable=no-self-use
90+
def instrument_connection(self, connection):
91+
"""Enable instrumentation in a Psycopg2 connection.
92+
93+
Args:
94+
connection: The connection to instrument.
95+
96+
Returns:
97+
An instrumented connection.
98+
"""
99+
tracer = get_tracer(__name__, __version__)
100+
101+
return dbapi.instrument_connection(
102+
tracer,
103+
connection,
104+
self._DATABASE_COMPONENT,
105+
self._DATABASE_TYPE,
106+
self._CONNECTION_ATTRIBUTES,
107+
)
108+
109+
def uninstrument_connection(self, connection):
110+
"""Disable instrumentation in a Psycopg2 connection.
111+
112+
Args:
113+
connection: The connection to uninstrument.
114+
115+
Returns:
116+
An uninstrumented connection.
117+
"""
118+
return dbapi.uninstrument_connection(connection)

0 commit comments

Comments
 (0)