Skip to content

Commit

Permalink
Update
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-evandenberg committed Sep 10, 2024
1 parent 9f51335 commit 11fab71
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/snowflake/snowpark/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -3399,7 +3399,7 @@ def flatten(
set_api_call_source(df, "Session.flatten")
return df

def query_history(self) -> QueryHistory:
def query_history(self, include_describe: bool = False) -> QueryHistory:
"""Create an instance of :class:`QueryHistory` as a context manager to record queries that are pushed down to the Snowflake database.
>>> with session.query_history() as query_history:
Expand All @@ -3408,7 +3408,7 @@ def query_history(self) -> QueryHistory:
... res = df.collect()
>>> assert len(query_history.queries) == 1
"""
query_listener = QueryHistory(self)
query_listener = QueryHistory(self, include_describe)
self._conn.add_query_listener(query_listener)
return query_listener

Expand Down
14 changes: 14 additions & 0 deletions tests/integ/test_query_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,20 @@ def test_query_history(session):
assert len(query_listener.queries) == 1
assert query_listener.queries[0].query_id is not None
assert query_listener.queries[0].sql_text == "select 0"
assert not query_listener.queries[0].is_describe


def test_query_history_with_describe(session):
with session.query_history(True) as query_listener:
df = session.sql("select 0")
df.columns
df.collect()
assert len(query_listener.queries) == 2
for query in query_listener.queries:
assert query.query_id is not None
assert query.sql_text == "select 0"
assert query_listener.queries[0].is_describe
assert not query_listener.queries[1].is_describe


def test_query_history_stop_listening(session):
Expand Down

0 comments on commit 11fab71

Please sign in to comment.