def test_signature_create_keyspace():
assert (
extract_signature(
"CREATE KEYSPACE testkeyspace WITH REPLICATION = { 'class' : 'NetworkTopologyStrategy', 'datacenter1' : 3 };"
)
== "CREATE KEYSPACE"
)
def call(self, module, method, wrapped, instance, args, kwargs):
name = self.get_wrapped_name(wrapped, instance, method)
context = {}
if method == "Cluster.connect":
span_action = "connect"
if hasattr(instance,
"contact_points_resolved"): # < cassandra-driver 3.18
host = instance.contact_points_resolved[0]
port = instance.port
else:
host = instance.endpoints_resolved[0].address
port = instance.endpoints_resolved[0].port
else:
hosts = list(instance.hosts)
if hasattr(hosts[0], "endpoint"):
host = hosts[0].endpoint.address
port = hosts[0].endpoint.port
else:
# < cassandra-driver 3.18
host = hosts[0].address
port = instance.cluster.port
span_action = "query"
query = args[0] if args else kwargs.get("query")
if hasattr(query, "query_string"):
query_str = query.query_string
elif hasattr(query, "prepared_statement") and hasattr(
query.prepared_statement, "query"):
query_str = query.prepared_statement.query
elif isinstance(query, compat.string_types):
query_str = query
else:
query_str = None
if query_str:
name = extract_signature(query_str)
context["db"] = {"type": "sql", "statement": query_str}
context["destination"] = {
"address": host,
"port": port,
"service": {
"name": "cassandra",
"resource": "cassandra",
"type": "db"
},
}
with capture_span(name,
span_type="db",
span_subtype="cassandra",
span_action=span_action,
extra=context):
return wrapped(*args, **kwargs)
def test_signature_create_columnfamily():
assert (
extract_signature(
"""CREATE COLUMNFAMILY users (
userid text PRIMARY KEY,
first_name text,
last_name text,
emails set<text>,
top_scores list<int>,
todo map<timestamp, text>
);"""
)
== "CREATE COLUMNFAMILY"
)
def _trace_sql(self, method, sql, params):
signature = extract_signature(sql)
with capture_span(
signature,
span_type="db",
span_subtype="sqlite",
span_action="query",
extra={"db": {
"type": "sql",
"statement": sql
}},
):
if params is None:
return method(sql)
else:
return method(sql, params)
async def call(self, module, method, wrapped, instance, args, kwargs):
if method == "Cursor.execute":
query = args[0] if len(args) else kwargs["operation"]
query = _bake_sql(instance.raw, query)
name = extract_signature(query)
context = {"db": {"type": "sql", "statement": query}}
action = "query"
elif method == "Cursor.callproc":
func = args[0] if len(args) else kwargs["procname"]
name = func + "()"
context = None
action = "exec"
else:
raise AssertionError("call from uninstrumented method")
async with async_capture_span(name,
leaf=True,
span_type="db",
span_subtype="postgres",
span_action=action,
extra=context):
return await wrapped(*args, **kwargs)
def extract_signature(self, sql):
return extract_signature(sql)
def test_select_from_collection():
assert extract_signature("SELECT first, last FROM a.b WHERE id = 1;") == "SELECT FROM a.b"