diff --git a/env/integtest_pg_conn.py b/env/integtest_pg_conn.py
index 8b63b24e..19afca25 100644
--- 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
index 3eb91780..cbc7e211 100644
--- 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"