def test_executemany(self):
db_integration = DatabaseApiIntegration(self.tracer, "testcomponent")
mock_connection = db_integration.wrapped_connection(
mock_connect, {}, {})
cursor = mock_connection.cursor()
cursor.executemany("Test query")
self.assertTrue(self.start_as_current_span.called)
self.assertEqual(self.span.attributes["db.statement"], "Test query")
def test_executemany(self):
db_integration = DatabaseApiIntegration(self.tracer, "testcomponent")
mock_connection = db_integration.wrapped_connection(
mock_connect, {}, {})
cursor = mock_connection.cursor()
cursor.executemany("Test query")
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
self.assertEqual(span.attributes["db.statement"], "Test query")
def test_span_failed(self):
db_integration = DatabaseApiIntegration(self.tracer, "testcomponent")
mock_connection = db_integration.wrapped_connection(
mock_connect, {}, {})
cursor = mock_connection.cursor()
try:
cursor.execute("Test query", throw_exception=True)
except Exception: # pylint: disable=broad-except
self.assertEqual(self.span.attributes["db.statement"],
"Test query")
self.assertIs(
self.span.status.canonical_code,
trace_api.status.StatusCanonicalCode.UNKNOWN,
)
self.assertEqual(self.span.status.description, "Test Exception")
def trace_integration(tracer_provider=None):
"""Integrate with PostgreSQL Psycopg library.
Psycopg: http://initd.org/psycopg/
"""
tracer = get_tracer(__name__, __version__, tracer_provider)
connection_attributes = {
"database": "info.dbname",
"port": "info.port",
"host": "info.host",
"user": "******",
}
db_integration = DatabaseApiIntegration(
tracer,
DATABASE_COMPONENT,
database_type=DATABASE_TYPE,
connection_attributes=connection_attributes,
)
# pylint: disable=unused-argument
def wrap_connect(
connect_func: typing.Callable[..., any],
instance: typing.Any,
args: typing.Tuple[any, any],
kwargs: typing.Dict[any, any],
):
connection = connect_func(*args, **kwargs)
db_integration.get_connection_attributes(connection)
connection.cursor_factory = PsycopgTraceCursor
return connection
try:
wrapt.wrap_function_wrapper(psycopg2, "connect", wrap_connect)
except Exception as ex: # pylint: disable=broad-except
logger.warning("Failed to integrate with pyscopg2. %s", str(ex))
class PsycopgTraceCursor(psycopg2.extensions.cursor):
def __init__(self, *args, **kwargs):
self._traced_cursor = TracedCursor(db_integration)
super(PsycopgTraceCursor, self).__init__(*args, **kwargs)
# pylint: disable=redefined-builtin
def execute(self, query, vars=None):
if isinstance(query, Composable):
query = query.as_string(self)
return self._traced_cursor.traced_execution(
super(PsycopgTraceCursor, self).execute, query, vars)
# pylint: disable=redefined-builtin
def executemany(self, query, vars):
if isinstance(query, Composable):
query = query.as_string(self)
return self._traced_cursor.traced_execution(
super(PsycopgTraceCursor, self).executemany, query, vars)
# pylint: disable=redefined-builtin
def callproc(self, procname, vars=None):
return self._traced_cursor.traced_execution(
super(PsycopgTraceCursor, self).callproc, procname, vars)
def test_span_failed(self):
db_integration = DatabaseApiIntegration(self.tracer, "testcomponent")
mock_connection = db_integration.wrapped_connection(
mock_connect, {}, {})
cursor = mock_connection.cursor()
with self.assertRaises(Exception):
cursor.execute("Test query", throw_exception=True)
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
self.assertEqual(span.attributes["db.statement"], "Test query")
self.assertIs(
span.status.canonical_code,
trace_api.status.StatusCanonicalCode.UNKNOWN,
)
self.assertEqual(span.status.description, "Test Exception")
def test_span_succeeded(self):
connection_props = {
"database": "testdatabase",
"server_host": "testhost",
"server_port": 123,
"user": "******",
}
connection_attributes = {
"database": "database",
"port": "server_port",
"host": "server_host",
"user": "******",
}
db_integration = DatabaseApiIntegration(self.tracer, "testcomponent",
"testtype",
connection_attributes)
mock_connection = db_integration.wrapped_connection(
mock_connect, {}, connection_props)
cursor = mock_connection.cursor()
cursor.execute("Test query", ("param1Value", False))
self.assertTrue(self.start_as_current_span.called)
self.assertEqual(
self.start_as_current_span.call_args[0][0],
"testcomponent.testdatabase",
)
self.assertIs(
self.start_as_current_span.call_args[1]["kind"],
trace_api.SpanKind.CLIENT,
)
self.assertEqual(self.span.attributes["component"], "testcomponent")
self.assertEqual(self.span.attributes["db.type"], "testtype")
self.assertEqual(self.span.attributes["db.instance"], "testdatabase")
self.assertEqual(self.span.attributes["db.statement"], "Test query")
self.assertEqual(
self.span.attributes["db.statement.parameters"],
"('param1Value', False)",
)
self.assertEqual(self.span.attributes["db.user"], "testuser")
self.assertEqual(self.span.attributes["net.peer.name"], "testhost")
self.assertEqual(self.span.attributes["net.peer.port"], 123)
self.assertIs(
self.span.status.canonical_code,
trace_api.status.StatusCanonicalCode.OK,
)
def test_span_succeeded(self):
connection_props = {
"database": "testdatabase",
"server_host": "testhost",
"server_port": 123,
"user": "******",
}
connection_attributes = {
"database": "database",
"port": "server_port",
"host": "server_host",
"user": "******",
}
db_integration = DatabaseApiIntegration(self.tracer, "testcomponent",
"testtype",
connection_attributes)
mock_connection = db_integration.wrapped_connection(
mock_connect, {}, connection_props)
cursor = mock_connection.cursor()
cursor.execute("Test query", ("param1Value", False))
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
self.assertEqual(span.name, "testcomponent.testdatabase")
self.assertIs(span.kind, trace_api.SpanKind.CLIENT)
self.assertEqual(span.attributes["component"], "testcomponent")
self.assertEqual(span.attributes["db.type"], "testtype")
self.assertEqual(span.attributes["db.instance"], "testdatabase")
self.assertEqual(span.attributes["db.statement"], "Test query")
self.assertEqual(
span.attributes["db.statement.parameters"],
"('param1Value', False)",
)
self.assertEqual(span.attributes["db.user"], "testuser")
self.assertEqual(span.attributes["net.peer.name"], "testhost")
self.assertEqual(span.attributes["net.peer.port"], 123)
self.assertIs(
span.status.canonical_code,
trace_api.status.StatusCanonicalCode.OK,
)