Include count in execute_write_many traces, closes #1571

This commit is contained in:
Simon Willison 2021-12-19 12:30:34 -08:00
commit f65817000f
3 changed files with 28 additions and 9 deletions

View file

@ -114,11 +114,22 @@ class Database:
async def execute_write_many(self, sql, params_seq, block=False):
def _inner(conn):
with conn:
return conn.executemany(sql, params_seq)
count = 0
with trace("sql", database=self.name, sql=sql.strip(), executemany=True):
results = await self.execute_write_fn(_inner, block=block)
def count_params(params):
nonlocal count
for param in params:
count += 1
yield param
with conn:
return conn.executemany(sql, count_params(params_seq)), count
with trace(
"sql", database=self.name, sql=sql.strip(), executemany=True
) as kwargs:
results, count = await self.execute_write_fn(_inner, block=block)
kwargs["count"] = count
return results
async def execute_write_fn(self, fn, block=False):

View file

@ -32,14 +32,14 @@ def trace(type, **kwargs):
), f".trace() keyword parameters cannot include {TRACE_RESERVED_KEYS}"
task_id = get_task_id()
if task_id is None:
yield
yield kwargs
return
tracer = tracers.get(task_id)
if tracer is None:
yield
yield kwargs
return
start = time.perf_counter()
yield
yield kwargs
end = time.perf_counter()
trace_info = {
"type": type,