wrote test for invalid hint (and fixed hint check notice)

Author: wangpatrick57

env/integtest_pg_conn.py

@@ -129,9 +129,7 @@ def test_time_query(self) -> None:
         self.assertIsNone(explain_data)
 
     def test_time_query_with_explain(self) -> None:
-        _, _, explain_data = self.pg_conn.time_query(
-            "select pg_sleep(1)", add_explain=True
-        )
+        _, _, explain_data = self.pg_conn.time_query("select 1", add_explain=True)
         self.assertIsNotNone(explain_data)
 
     def test_time_query_with_timeout(self) -> None:
@@ -150,7 +148,7 @@ def test_time_query_with_invalid_table(self) -> None:
         with self.assertRaises(psycopg.errors.UndefinedTable):
             self.pg_conn.time_query("select * from itemline limit 10")
 
-    def test_time_query_with_hint(self) -> None:
+    def test_time_query_with_valid_hints(self) -> None:
         join_query = """SELECT *
 FROM orders
 JOIN lineitem ON o_orderkey = l_orderkey
@@ -168,9 +166,17 @@ def test_time_query_with_hint(self) -> None:
                 query_knobs=[f"{hint_join_type}(lineitem orders)"],
                 add_explain=True,
             )
+            assert explain_data is not None  # This assertion is for mypy.
             actual_join_type = explain_data["Plan"]["Plans"][0]["Node Type"]
             self.assertEqual(expected_join_type, actual_join_type)
 
+    def test_time_query_with_invalid_hint(self) -> None:
+        with self.assertRaises(RuntimeError) as context:
+            self.pg_conn.time_query("select 1", query_knobs=["dbgym"])
+            self.assertTrue(
+                'Unrecognized hint keyword "dbgym"' in str(context.exception)
+            )
+
 
 if __name__ == "__main__":
     unittest.main()

env/pg_conn.py

@@ -70,6 +70,7 @@ def __init__(
         self.dbdata_dpath = self.dbdata_parent_dpath / f"dbdata{self.pgport}"
 
         self._conn: Optional[psycopg.Connection[Any]] = None
+        self.hint_check_failed_with: Optional[str] = None
 
     def get_kv_connstr(self) -> str:
         return get_kv_connstr(self.pgport)
@@ -79,6 +80,21 @@ def conn(self) -> psycopg.Connection[Any]:
             self._conn = psycopg.connect(
                 self.get_kv_connstr(), autocommit=True, prepare_threshold=None
             )
+
+            def hint_check_notice_handler(notice: psycopg.errors.Diagnostic) -> None:
+                """
+                Custom handler for raising errors if hints fail.
+                """
+                if (
+                    notice.message_detail is not None
+                    and "hint" in notice.message_detail.lower()
+                ):
+                    self.hint_check_failed_with = notice.message_detail
+
+            # We add the notice handler when the _conn is created instead of before executing a
+            # query to avoid adding it more than once.
+            self._conn.add_notice_handler(hint_check_notice_handler)
+
         return self._conn
 
     def disconnect(self) -> None:
@@ -137,17 +153,6 @@ def time_query(
         did_time_out = False
         explain_data = None
 
-        # def hint_notice_handler(notice) -> None:
-        #     """
-        #     Custom handler for database notices.
-        #     Raises an error or logs the notice if it indicates a problem.
-        #     """
-        #     logging.getLogger(DBGYM_LOGGER_NAME).warning(f"Postgres notice: {notice}")
-        #     if "hint" in notice.message.lower():
-        #         raise RuntimeError(f"Query hint failed: {notice.message}")
-
-        # self.conn().add_notice_handler(hint_notice_handler)
-
         try:
             if query_knobs:
                 query = f"/*+ {' '.join(query_knobs)} */ {query}"
@@ -158,10 +163,16 @@ def time_query(
                 ), "If you're using add_explain, don't also write explain manually in the query."
                 query = f"explain (analyze, format json, timing off) {query}"
 
+            # Reset this every time before calling execute() so that hint_check_notice_handler works correctly.
+            self.hint_check_failed_with = None
+
             start_time = time.time()
             cursor = self.conn().execute(query)
             qid_runtime = (time.time() - start_time) * 1e6
 
+            if self.hint_check_failed_with is not None:
+                raise RuntimeError(f"Query hint failed: {self.hint_check_failed_with}")
+
             if add_explain:
                 c = [c for c in cursor][0][0][0]
                 assert "Execution Time" in c