def test_commit_is_traced(self):
connection = self.connection
tracer = self.tracer
connection.commit.return_value = None
pin = Pin('pin_name', tracer=tracer)
traced_connection = TracedConnection(connection, pin)
traced_connection.commit()
assert tracer.pop()[0].name == 'mock.connection.commit'
connection.commit.assert_called_with()
def test_rollback_is_traced(self):
connection = self.connection
tracer = self.tracer
connection.rollback.return_value = None
pin = Pin('pin_name', tracer=tracer)
traced_connection = TracedConnection(connection, pin)
traced_connection.rollback()
assert tracer.writer.pop()[0].name == 'mock.connection.rollback'
connection.rollback.assert_called_with()
def test_connection_analytics_with_rate(self):
with self.override_config("dbapi2", dict(analytics_enabled=True, analytics_sample_rate=0.5)):
connection = self.connection
tracer = self.tracer
connection.commit.return_value = None
pin = Pin("pin_name", tracer=tracer)
traced_connection = TracedConnection(connection, pin)
traced_connection.commit()
span = tracer.pop()[0]
self.assertIsNone(span.get_metric(ANALYTICS_SAMPLE_RATE_KEY))
def test_cursor_class(self):
pin = Pin('pin_name', tracer=self.tracer)
# Default
traced_connection = TracedConnection(self.connection, pin=pin)
self.assertTrue(traced_connection._self_cursor_cls is TracedCursor)
# Trace fetched methods
with self.override_config('dbapi2', dict(trace_fetch_methods=True)):
traced_connection = TracedConnection(self.connection, pin=pin)
self.assertTrue(traced_connection._self_cursor_cls is FetchTracedCursor)
# Manually provided cursor class
with self.override_config('dbapi2', dict(trace_fetch_methods=True)):
traced_connection = TracedConnection(self.connection, pin=pin, cursor_cls=TracedCursor)
self.assertTrue(traced_connection._self_cursor_cls is TracedCursor)
def patch_conn(conn):
tags = {t: getattr(conn, a, '') for t, a in CONN_ATTR_BY_TAG.items()}
pin = Pin(app='pymysql', tags=tags)
# grab the metadata from the conn
wrapped = TracedConnection(conn, pin=pin, cfg=config.pymysql)
pin.onto(wrapped)
return wrapped
def patch_conn(conn):
tags = {t: getattr(conn, a, '') for t, a in CONN_ATTR_BY_TAG.items()}
pin = Pin(service="mysql", app="mysql", app_type="db", tags=tags)
# grab the metadata from the conn
wrapped = TracedConnection(conn)
pin.onto(wrapped)
return wrapped
def patch_conn(conn):
tags = {t: getattr(conn, a) for t, a in CONN_ATTR_BY_TAG.items() if getattr(conn, a, '') != ''}
pin = Pin(service='mysql', app='mysql', app_type=AppTypes.db, tags=tags)
# grab the metadata from the conn
wrapped = TracedConnection(conn, pin=pin)
pin.onto(wrapped)
return wrapped
def patch_conn(conn, *args, **kwargs):
tags = {
t: kwargs[k] if k in kwargs else args[p] for t, (k, p) in KWPOS_BY_TAG.items() if k in kwargs or len(args) > p
}
tags[net.TARGET_PORT] = conn.port
pin = Pin(app="mysql", tags=tags)
# grab the metadata from the conn
wrapped = TracedConnection(conn, pin=pin, cfg=config.mysqldb)
pin.onto(wrapped)
return wrapped
def patch_conn(conn, *args, **kwargs):
tags = {
t: kwargs[k] if k in kwargs else args[p]
for t, (k, p) in KWPOS_BY_TAG.items() if k in kwargs or len(args) > p
}
tags[net.TARGET_PORT] = conn.port
pin = Pin(service='mysql', app='mysql', app_type=AppTypes.db, tags=tags)
# grab the metadata from the conn
wrapped = TracedConnection(conn, pin=pin)
pin.onto(wrapped)
return wrapped
def _connect(func, instance, args, kwargs):
conn = func(*args, **kwargs)
tags = {
net.TARGET_HOST: kwargs["host"],
net.TARGET_PORT: kwargs["port"],
db.USER: kwargs["user"],
db.NAME: kwargs["database"],
}
pin = Pin(app="mariadb", tags=tags)
wrapped = TracedConnection(conn, pin=pin, cfg=config.mariadb)
pin.onto(wrapped)
return wrapped
def test_cursor_class(self):
pin = Pin("pin_name", tracer=self.tracer)
# Default
traced_connection = TracedConnection(self.connection, pin=pin)
self.assertTrue(traced_connection._self_cursor_cls is TracedCursor)
# Trace fetched methods
with self.override_config("dbapi2", dict(trace_fetch_methods=True)):
with warnings.catch_warnings(record=True) as ws:
warnings.simplefilter("always")
traced_connection = TracedConnection(self.connection, pin=pin)
self.assertTrue(traced_connection._self_cursor_cls is FetchTracedCursor)
(w,) = ws
self.assertTrue(issubclass(w.category, DeprecationWarning))
self.assertTrue("ddtrace.config.dbapi2.trace_fetch_methods is now deprecated" in str(w.message))
# Manually provided cursor class
with self.override_config("dbapi2", dict(trace_fetch_methods=True)):
traced_connection = TracedConnection(self.connection, pin=pin, cursor_cls=TracedCursor)
self.assertTrue(traced_connection._self_cursor_cls is TracedCursor)
def test_connection_context_manager(self):
class Cursor(object):
rowcount = 0
def execute(self, *args, **kwargs):
pass
def fetchall(self, *args, **kwargs):
pass
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def commit(self, *args, **kwargs):
pass
# When a connection is returned from a context manager the object proxy
# should be returned so that tracing works.
class ConnectionConnection(object):
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def cursor(self):
return Cursor()
def commit(self):
pass
pin = Pin("pin", tracer=self.tracer)
conn = TracedConnection(ConnectionConnection(), pin)
with conn as conn2:
conn2.commit()
spans = self.tracer.writer.pop()
assert len(spans) == 1
with conn as conn2:
with conn2.cursor() as cursor:
cursor.execute("query")
cursor.fetchall()
spans = self.tracer.writer.pop()
assert len(spans) == 1
# If a cursor is returned from the context manager
# then it should be instrumented.
class ConnectionCursor(object):
def __enter__(self):
return Cursor()
def __exit__(self, *exc):
return False
def commit(self):
pass
with TracedConnection(ConnectionCursor(), pin) as cursor:
cursor.execute("query")
cursor.fetchall()
spans = self.tracer.writer.pop()
assert len(spans) == 1
# If a traced cursor is returned then it should not
# be double instrumented.
class ConnectionTracedCursor(object):
def __enter__(self):
return self.cursor()
def __exit__(self, *exc):
return False
def cursor(self):
return TracedCursor(Cursor(), pin, {})
def commit(self):
pass
with TracedConnection(ConnectionTracedCursor(), pin) as cursor:
cursor.execute("query")
cursor.fetchall()
spans = self.tracer.writer.pop()
assert len(spans) == 1
# Check when a different connection object is returned
# from a connection context manager.
# No traces should be produced.
other_conn = ConnectionConnection()
class ConnectionDifferentConnection(object):
def __enter__(self):
return other_conn
def __exit__(self, *exc):
return False
def cursor(self):
return Cursor()
def commit(self):
pass
conn = TracedConnection(ConnectionDifferentConnection(), pin)
with conn as conn2:
conn2.commit()
spans = self.tracer.writer.pop()
assert len(spans) == 0
with conn as conn2:
with conn2.cursor() as cursor:
cursor.execute("query")
cursor.fetchall()
spans = self.tracer.writer.pop()
assert len(spans) == 0
# When some unexpected value is returned from the context manager
# it should be handled gracefully.
class ConnectionUnknown(object):
def __enter__(self):
return 123456
def __exit__(self, *exc):
return False
def cursor(self):
return Cursor()
def commit(self):
pass
conn = TracedConnection(ConnectionDifferentConnection(), pin)
with conn as conn2:
conn2.commit()
spans = self.tracer.writer.pop()
assert len(spans) == 0
with conn as conn2:
with conn2.cursor() as cursor:
cursor.execute("query")
cursor.fetchall()
spans = self.tracer.writer.pop()
assert len(spans) == 0
# Errors should be the same when no context management is defined.
class ConnectionNoCtx(object):
def cursor(self):
return Cursor()
def commit(self):
pass
conn = TracedConnection(ConnectionNoCtx(), pin)
with pytest.raises(AttributeError):
with conn:
pass
with pytest.raises(AttributeError):
with conn as conn2:
pass
spans = self.tracer.writer.pop()
assert len(spans) == 0
