From 875ad1a15c61a435c41e062aba4848aca247203e Mon Sep 17 00:00:00 2001
From: Patrick Wang <wang.patrick57@gmail.com>
Date: Wed, 13 Nov 2024 19:59:28 -0500
Subject: [PATCH] removed prefix from time_query()

---
 env/pg_conn.py                  | 11 ++++++++---
 tune/protox/env/util/execute.py |  4 +---
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/env/pg_conn.py b/env/pg_conn.py
index a91bc36f..15f0e7e7 100644
--- a/env/pg_conn.py
+++ b/env/pg_conn.py
@@ -100,8 +100,13 @@ def move_log(self) -> None:
             self.log_step += 1
 
     def time_query(
-        self, prefix: str, query: str, timeout: float
+        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.
+
+        It returns the runtime, whether the query timed out, and the explain data.
+        """
         did_time_out = False
         has_explain = "EXPLAIN" in query
         explain_data = None
@@ -118,12 +123,12 @@ def time_query(
                 explain_data = c
 
             logging.getLogger(DBGYM_LOGGER_NAME).debug(
-                f"{prefix} evaluated in {qid_runtime/1e6}"
+                f"{query} evaluated in {qid_runtime/1e6}"
             )
 
         except QueryCanceled:
             logging.getLogger(DBGYM_LOGGER_NAME).debug(
-                f"{prefix} exceeded evaluation timeout {timeout}"
+                f"{query} exceeded evaluation timeout {timeout}"
             )
             qid_runtime = timeout * 1e6
             did_time_out = True
diff --git a/tune/protox/env/util/execute.py b/tune/protox/env/util/execute.py
index 130b60ae..b11acff8 100644
--- a/tune/protox/env/util/execute.py
+++ b/tune/protox/env/util/execute.py
@@ -34,7 +34,6 @@ def _force_statement_timeout(
 
 
 def _acquire_metrics_around_query(
-    prefix: str,
     pg_conn: PostgresConn,
     query: str,
     query_timeout: float = 0.0,
@@ -52,7 +51,7 @@ def _acquire_metrics_around_query(
         ), 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(
-        prefix, query, query_timeout
+        query, query_timeout
     )
 
     # Wipe the statement timeout.
@@ -108,7 +107,6 @@ def execute_variations(
         )
 
         runtime, did_time_out, explain_data, metric = _acquire_metrics_around_query(
-            prefix=qr.prefix_qid,
             pg_conn=pg_conn,
             query=pqk_query,
             query_timeout=timeout_limit,