Skip to content

Commit

Permalink
Time all db calls (#643)
Browse files Browse the repository at this point in the history
  • Loading branch information
cmyui authored Feb 28, 2024
1 parent 6b64b8b commit fa94c6d
Show file tree
Hide file tree
Showing 2 changed files with 93 additions and 7 deletions.
73 changes: 66 additions & 7 deletions app/adapters/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

from app import settings
from app.logging import log
from app.timer import Timer


class MySQLDialect(MySQLDialect_mysqldb):
Expand Down Expand Up @@ -49,11 +50,18 @@ async def fetch_one(
if isinstance(query, ClauseElement):
query, params = self._compile(query)

row = await self._database.fetch_one(query, params)
with Timer() as timer:
row = await self._database.fetch_one(query, params)

if settings.DEBUG:
time_elapsed = timer.elapsed()
log(
f"Executed SQL query: {query} {params}",
extra={"query": query, "params": params},
f"Executed SQL query: {query} {params} in {time_elapsed:.2f} seconds.",
extra={
"query": query,
"params": params,
"time_elapsed": time_elapsed,
},
)

return dict(row._mapping) if row is not None else None
Expand All @@ -66,7 +74,20 @@ async def fetch_all(
if isinstance(query, ClauseElement):
query, params = self._compile(query)

rows = await self._database.fetch_all(query, params)
with Timer() as timer:
rows = await self._database.fetch_all(query, params)

if settings.DEBUG:
time_elapsed = timer.elapsed()
log(
f"Executed SQL query: {query} {params} in {time_elapsed:.2f} seconds.",
extra={
"query": query,
"params": params,
"time_elapsed": time_elapsed,
},
)

return [dict(row._mapping) for row in rows]

async def fetch_val(
Expand All @@ -78,14 +99,40 @@ async def fetch_val(
if isinstance(query, ClauseElement):
query, params = self._compile(query)

val = await self._database.fetch_val(query, params, column)
with Timer() as timer:
val = await self._database.fetch_val(query, params, column)

if settings.DEBUG:
time_elapsed = timer.elapsed()
log(
f"Executed SQL query: {query} {params} in {time_elapsed:.2f} seconds.",
extra={
"query": query,
"params": params,
"time_elapsed": time_elapsed,
},
)

return val

async def execute(self, query: MySQLQuery, params: MySQLParams = None) -> int:
if isinstance(query, ClauseElement):
query, params = self._compile(query)

rec_id = await self._database.execute(query, params)
with Timer() as timer:
rec_id = await self._database.execute(query, params)

if settings.DEBUG:
time_elapsed = timer.elapsed()
log(
f"Executed SQL query: {query} {params} in {time_elapsed:.2f} seconds.",
extra={
"query": query,
"params": params,
"time_elapsed": time_elapsed,
},
)

return cast(int, rec_id)

# NOTE: this accepts str since current execute_many uses are not using alchemy.
Expand All @@ -94,7 +141,19 @@ async def execute_many(self, query: str, params: list[MySQLParams]) -> None:
if isinstance(query, ClauseElement):
query, _ = self._compile(query)

await self._database.execute_many(query, params)
with Timer() as timer:
await self._database.execute_many(query, params)

if settings.DEBUG:
time_elapsed = timer.elapsed()
log(
f"Executed SQL query: {query} {params} in {time_elapsed:.2f} seconds.",
extra={
"query": query,
"params": params,
"time_elapsed": time_elapsed,
},
)

def transaction(
self,
Expand Down
27 changes: 27 additions & 0 deletions app/timer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
from __future__ import annotations

import time
from types import TracebackType


class Timer:
def __init__(self) -> None:
self.start_time: float | None = None
self.end_time: float | None = None

def __enter__(self) -> Timer:
self.start_time = time.time()
return self

def __exit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
traceback: TracebackType | None,
) -> None:
self.end_time = time.time()

def elapsed(self) -> float:
if self.start_time is None or self.end_time is None:
raise ValueError("Timer has not been started or stopped.")
return self.end_time - self.start_time

0 comments on commit fa94c6d

Please sign in to comment.