changed return type of time_query

cmu-db · Nov 14, 2024 · 3d714e8 · 3d714e8
1 parent 31f782d
commit 3d714e8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/env/integtest_pg_conn.py b/env/integtest_pg_conn.py
@@ -7,7 +7,11 @@
 from psycopg.errors import QueryCanceled
 
 from env.pg_conn import PostgresConn
-from util.pg import DEFAULT_POSTGRES_PORT, get_is_postgres_running, get_running_postgres_ports
+from util.pg import (
+    DEFAULT_POSTGRES_PORT,
+    get_is_postgres_running,
+    get_running_postgres_ports,
+)
 from util.workspace import (
     DEFAULT_BOOT_CONFIG_FPATH,
     DBGymConfig,

diff --git a/env/pg_conn.py b/env/pg_conn.py
@@ -109,7 +109,9 @@ def force_statement_timeout(self, timeout_sec: float) -> None:
             except QueryCanceled:
                 retry = True
 
-    def time_query(self, query: str, timeout_sec: float = 0) -> tuple[float, bool, Any]:
+    def time_query(
+        self, query: str, timeout_sec: float = 0
+    ) -> tuple[float, bool, Optional[dict[str, Any]]]:
         """
         Run a query with a timeout. If you want to attach per-query knobs, attach them to the query string itself.
         Following Postgres's convention, timeout=0 indicates "disable timeout"