def test_instrument_connection(self):
connection = mock.Mock()
# Avoid get_attributes failing because can't concatenate mock
connection.database = "-"
connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
self.assertIsInstance(connection2, dbapi.TracedConnectionProxy)
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 = dbapi.instrument_connection(self.tracer, connection, "-")
self.assertIs(connection2.__wrapped__, connection)
connection3 = dbapi.uninstrument_connection(connection2)
self.assertIs(connection3, connection)
with self.assertLogs(level=logging.WARNING):
connection4 = dbapi.uninstrument_connection(connection)
self.assertIs(connection4, connection)
def instrument_connection(self, connection):
"""Enable instrumentation in a PyMySQL connection.
Args:
connection: The connection to instrument.
Returns:
An instrumented connection.
"""
tracer = get_tracer(__name__, __version__)
return dbapi.instrument_connection(
tracer,
connection,
self._DATABASE_COMPONENT,
self._DATABASE_TYPE,
self._CONNECTION_ATTRIBUTES,
)