def test_instrument_connection(self):
connection = mock.Mock()
# Avoid get_attributes failing because can't concatenate mock
connection.database = "-"
connection2 = wrappers.instrument_connection(self.tracer, connection,
"-")
self.assertIs(connection2.__wrapped__, connection)
def test_uninstrument_connection(self):
connection = mock.Mock()
# Set connection.database to avoid a failure because mock can't
# be concatenated
connection.database = "-"
connection2 = wrappers.instrument_connection(self.tracer, connection,
"-")
self.assertIs(connection2.__wrapped__, connection)
connection3 = wrappers.uninstrument_connection(connection2)
self.assertIs(connection3, connection)
with self.assertLogs(level=logging.WARNING):
connection4 = wrappers.uninstrument_connection(connection)
self.assertIs(connection4, connection)
def instrument_connection(self, connection):
"""Enable instrumentation in a aiopg connection.
Args:
connection: The connection to instrument.
Returns:
An instrumented connection.
"""
return wrappers.instrument_connection(
__name__,
connection,
self._DATABASE_SYSTEM,
self._CONNECTION_ATTRIBUTES,
)
def instrument_connection(self, connection, tracer_provider=None):
"""Enable instrumentation in a aiopg connection.
Args:
connection: The connection to instrument.
tracer_provider: The optional tracer provider to use. If omitted
the current globally configured one is used.
Returns:
An instrumented connection.
"""
return wrappers.instrument_connection(
__name__,
connection,
self._DATABASE_SYSTEM,
self._CONNECTION_ATTRIBUTES,
version=__version__,
tracer_provider=tracer_provider,
)
