Skip to content

Commit

Permalink
wrote test_time_query_with_hint
Browse files Browse the repository at this point in the history
  • Loading branch information
wangpatrick57 committed Dec 21, 2024
1 parent 8455ebe commit 0cbcbba
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 22 deletions.
37 changes: 25 additions & 12 deletions env/integtest_pg_conn.py
Original file line number Diff line number Diff line change
Expand Up @@ -129,34 +129,47 @@ def test_time_query(self) -> None:
self.assertIsNone(explain_data)

def test_time_query_with_explain(self) -> None:
runtime, did_time_out, explain_data = self.pg_conn.time_query(
_, _, explain_data = self.pg_conn.time_query(
"select pg_sleep(1)", add_explain=True
)
self.assertTrue(abs(runtime - 1_000_000) < 100_000)
self.assertFalse(did_time_out)
self.assertIsNotNone(explain_data)

def test_time_query_with_timeout(self) -> None:
runtime, did_time_out, explain_data = self.pg_conn.time_query(
runtime, did_time_out, _ = self.pg_conn.time_query(
"select pg_sleep(3)", timeout=2
)
# The runtime should be about what the timeout is.
self.assertTrue(abs(runtime - 2_000_000) < 100_000)
self.assertTrue(did_time_out)
self.assertIsNone(explain_data)

def test_time_query_with_valid_table(self) -> None:
_, did_time_out, explain_data = self.pg_conn.time_query(
"select * from lineitem limit 10"
)
self.assertFalse(did_time_out)
self.assertIsNone(explain_data)
# This just ensures that it doesn't raise any errors.
self.pg_conn.time_query("select * from lineitem limit 10")

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"
self.pg_conn.time_query("select * from itemline limit 10")

def test_time_query_with_hint(self) -> None:
join_query = """SELECT *
FROM orders
JOIN lineitem ON o_orderkey = l_orderkey
WHERE o_orderdate BETWEEN '1995-01-01' AND '1995-12-31'
LIMIT 10"""
join_types = [
("MergeJoin", "Merge Join"),
("HashJoin", "Hash Join"),
("NestLoop", "Nested Loop"),
]

for hint_join_type, expected_join_type in join_types:
_, _, explain_data = self.pg_conn.time_query(
join_query,
query_knobs=[f"{hint_join_type}(lineitem orders)"],
add_explain=True,
)
actual_join_type = explain_data["Plan"]["Plans"][0]["Node Type"]
self.assertEqual(expected_join_type, actual_join_type)


if __name__ == "__main__":
Expand Down
20 changes: 10 additions & 10 deletions env/pg_conn.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,16 +137,16 @@ 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)
# 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:
Expand Down

0 comments on commit 0cbcbba

Please sign in to comment.