diff --git a/env/pg_conn.py b/env/pg_conn.py index 15f0e7e7..ba1007e1 100644 --- a/env/pg_conn.py +++ b/env/pg_conn.py @@ -99,9 +99,7 @@ def move_log(self) -> None: shutil.move(pglog_fpath, pglog_this_step_fpath) self.log_step += 1 - def time_query( - self, query: str, timeout: float - ) -> tuple[float, bool, Any]: + def time_query(self, query: str, timeout: float) -> tuple[float, bool, Any]: """ Run a query with a timeout. If you want to attach per-query knobs, attach them to the query string itself. diff --git a/tune/protox/env/util/execute.py b/tune/protox/env/util/execute.py index b11acff8..45bbf417 100644 --- a/tune/protox/env/util/execute.py +++ b/tune/protox/env/util/execute.py @@ -50,9 +50,7 @@ def _acquire_metrics_around_query( query_timeout == 0 ), f'Setting query_timeout to 0 indicates "timeout". However, setting query_timeout ({query_timeout}) < 0 is a bug.' - qid_runtime, did_time_out, explain_data = pg_conn.time_query( - query, query_timeout - ) + qid_runtime, did_time_out, explain_data = pg_conn.time_query(query, query_timeout) # Wipe the statement timeout. _force_statement_timeout(pg_conn.conn(), 0)