def instrument_connection(
name: str,
connection,
database_component: str,
database_type: str = "",
connection_attributes: typing.Dict = None,
version: str = "",
tracer_provider: typing.Optional[TracerProvider] = None,
):
"""Enable instrumentation in a database connection.
Args:
name: Name of opentelemetry extension for aiopg.
connection: The connection to instrument.
database_component: Database driver name or database name "postgreSQL".
database_type: The Database type. For any SQL database, "sql".
connection_attributes: Attribute names for database, port, host and
user in a connection object.
version: Version of opentelemetry extension for aiopg.
tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
use. If ommited the current configured one is used.
Returns:
An instrumented connection.
"""
db_integration = AiopgIntegration(
name,
database_component,
database_type,
connection_attributes=connection_attributes,
version=version,
tracer_provider=tracer_provider,
)
db_integration.get_connection_attributes(connection)
return get_traced_connection_proxy(connection, db_integration)
def test_span_not_recording(self):
connection_props = {
"database": "testdatabase",
"server_host": "testhost",
"server_port": 123,
"user": "******",
}
connection_attributes = {
"database": "database",
"port": "server_port",
"host": "server_host",
"user": "******",
}
mock_tracer = mock.Mock()
mock_span = mock.Mock()
mock_span.is_recording.return_value = False
mock_tracer.start_span.return_value = mock_span
mock_tracer.use_span.return_value.__enter__ = mock_span
mock_tracer.use_span.return_value.__exit__ = True
db_integration = AiopgIntegration(
mock_tracer, "testcomponent", "testtype", connection_attributes
)
mock_connection = async_call(
db_integration.wrapped_connection(
mock_connect, {}, connection_props
)
)
cursor = async_call(mock_connection.cursor())
async_call(cursor.execute("Test query", ("param1Value", False)))
self.assertFalse(mock_span.is_recording())
self.assertTrue(mock_span.is_recording.called)
self.assertFalse(mock_span.set_attribute.called)
self.assertFalse(mock_span.set_status.called)
def test_executemany(self):
db_integration = AiopgIntegration(self.tracer, "testcomponent")
mock_connection = async_call(
db_integration.wrapped_connection(mock_connect, {}, {}))
cursor = async_call(mock_connection.cursor())
async_call(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_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 = AiopgIntegration(
self.tracer,
"testcomponent",
connection_attributes,
capture_parameters=True,
)
mock_connection = async_call(
db_integration.wrapped_connection(
mock_connect, {}, connection_props
)
)
cursor = async_call(mock_connection.cursor())
async_call(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, "Test")
self.assertIs(span.kind, trace_api.SpanKind.CLIENT)
self.assertEqual(
span.attributes[SpanAttributes.DB_SYSTEM], "testcomponent"
)
self.assertEqual(
span.attributes[SpanAttributes.DB_NAME], "testdatabase"
)
self.assertEqual(
span.attributes[SpanAttributes.DB_STATEMENT], "Test query"
)
self.assertEqual(
span.attributes["db.statement.parameters"],
"('param1Value', False)",
)
self.assertEqual(span.attributes[SpanAttributes.DB_USER], "testuser")
self.assertEqual(
span.attributes[SpanAttributes.NET_PEER_NAME], "testhost"
)
self.assertEqual(span.attributes[SpanAttributes.NET_PEER_PORT], 123)
self.assertIs(span.status.status_code, trace_api.StatusCode.UNSET)
def test_span_failed(self):
db_integration = AiopgIntegration(self.tracer, "testcomponent")
mock_connection = async_call(
db_integration.wrapped_connection(mock_connect, {}, {}))
cursor = async_call(mock_connection.cursor())
with self.assertRaises(Exception):
async_call(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.status_code, trace_api.StatusCode.ERROR)
self.assertEqual(span.status.description, "Exception: Test Exception")
def test_callproc(self):
db_integration = AiopgIntegration(self.tracer, "testcomponent")
mock_connection = async_call(
db_integration.wrapped_connection(mock_connect, {}, {}))
cursor = async_call(mock_connection.cursor())
async_call(cursor.callproc("Test stored procedure"))
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
self.assertEqual(
span.attributes[SpanAttributes.DB_STATEMENT],
"Test stored procedure",
)
def wrap_connect_(
wrapped: typing.Callable[..., typing.Any],
instance: typing.Any,
args: typing.Tuple[typing.Any, typing.Any],
kwargs: typing.Dict[typing.Any, typing.Any],
):
db_integration = AiopgIntegration(
name,
database_system,
connection_attributes=connection_attributes,
version=version,
tracer_provider=tracer_provider,
)
return _ContextManager( # pylint: disable=no-value-for-parameter
db_integration.wrapped_connection(wrapped, args, kwargs))
def wrap_create_pool_(
wrapped: typing.Callable[..., typing.Any],
instance: typing.Any,
args: typing.Tuple[typing.Any, typing.Any],
kwargs: typing.Dict[typing.Any, typing.Any],
):
db_integration = AiopgIntegration(
name,
database_system,
connection_attributes=connection_attributes,
version=version,
tracer_provider=tracer_provider,
)
return _PoolContextManager(
db_integration.wrapped_pool(wrapped, args, kwargs))
def wrap_connect_(
wrapped: typing.Callable[..., typing.Any],
instance: typing.Any,
args: typing.Tuple[typing.Any, typing.Any],
kwargs: typing.Dict[typing.Any, typing.Any],
):
db_integration = AiopgIntegration(
name,
database_component,
database_type=database_type,
connection_attributes=connection_attributes,
version=version,
tracer_provider=tracer_provider,
)
return _ContextManager(
db_integration.wrapped_connection(wrapped, args, kwargs))
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 = AiopgIntegration(
self.tracer, "testcomponent", "testtype", connection_attributes
)
mock_connection = async_call(
db_integration.wrapped_connection(
mock_connect, {}, connection_props
)
)
cursor = async_call(mock_connection.cursor())
async_call(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,
)
def instrument_connection(
name: str,
connection,
database_system: str,
connection_attributes: typing.Dict = None,
version: str = "",
tracer_provider: typing.Optional[TracerProvider] = None,
):
"""Enable instrumentation in a database connection.
Args:
name: Name of opentelemetry extension for aiopg.
connection: The connection to instrument.
database_system: An identifier for the database management system (DBMS)
product being used.
connection_attributes: Attribute names for database, port, host and
user in a connection object.
version: Version of opentelemetry extension for aiopg.
tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
use. If omitted the current configured one is used.
Returns:
An instrumented connection.
"""
if isinstance(connection, AsyncProxyObject):
logger.warning("Connection already instrumented")
return connection
db_integration = AiopgIntegration(
name,
database_system,
connection_attributes=connection_attributes,
version=version,
tracer_provider=tracer_provider,
)
db_integration.get_connection_attributes(connection)
return get_traced_connection_proxy(connection, db_integration)
async def wrap_create_pool_(
wrapped: typing.Callable[..., typing.Any],
instance: typing.Any,
args: typing.Tuple[typing.Any, typing.Any],
kwargs: typing.Dict[typing.Any, typing.Any],
):
db_integration = AiopgIntegration(
name,
database_component,
database_type,
connection_attributes=connection_attributes,
version=version,
tracer_provider=tracer_provider,
)
return await db_integration.wrapped_pool(wrapped, args, kwargs)
