SNOW-1557539 count describe queries for test (#2249)

<!---
Please answer these questions before creating your pull request. Thanks!
--->

1. Which Jira issue is this PR addressing? Make sure that there is an
accompanying issue to your PR.

   Fixes SNOW-1557539 count describe queries for test

2. Fill out the following pre-review checklist:

- [ ] I am adding a new automated test(s) to verify correctness of my
new code
- [ ] If this test skips Local Testing mode, I'm requesting review from
@snowflakedb/local-testing
   - [ ] I am adding new logging messages
   - [ ] I am adding a new telemetry message
   - [ ] I am adding new credentials
   - [ ] I am adding a new dependency
- [ ] If this is a new feature/behavior, I'm adding the Local Testing
parity changes.

3. Please describe how your code solves the related issue.

Add describe query to query listener history (opt-in)

Author: sfc-gh-evandenberg

CHANGELOG.md

@@ -13,6 +13,7 @@
 - Added support for `TimedeltaIndex.mean` method.
 - Added support for some cases of aggregating `Timedelta` columns on `axis=0` with `agg` or `aggregate`.
 - Added support for `by`, `left_by`, `right_by`, `left_index`, and `right_index` for `pd.merge_asof`.
+- Added support for passing parameter `include_describe` to `Session.query_history`.
 
 #### Bug Fixes
 

src/snowflake/snowpark/_internal/server_connection.py

@@ -244,9 +244,23 @@ def _get_string_datum(self, query: str) -> Optional[str]:
     @SnowflakePlan.Decorator.wrap_exception
     def get_result_attributes(self, query: str) -> List[Attribute]:
         return convert_result_meta_to_attribute(
-            run_new_describe(self._cursor, query), self.max_string_size
+            self._run_new_describe(self._cursor, query), self.max_string_size
         )
 
+    def _run_new_describe(
+        self, cursor: SnowflakeCursor, query: str
+    ) -> Union[List[ResultMetadata], List["ResultMetadataV2"]]:
+        result_metadata = run_new_describe(cursor, query)
+
+        for listener in filter(
+            lambda listener: hasattr(listener, "include_describe")
+            and listener.include_describe,
+            self._query_listener,
+        ):
+            listener._add_query(QueryRecord(cursor.sfqid, query, True))
+
+        return result_metadata
+
     @_Decorator.log_msg_and_perf_telemetry("Uploading file to stage")
     def upload_file(
         self,

src/snowflake/snowpark/query_history.py

@@ -12,6 +12,7 @@ class QueryRecord(NamedTuple):
 
     query_id: str
     sql_text: str
+    is_describe: bool = False
 
 
 class QueryHistory:
@@ -21,9 +22,14 @@ class QueryHistory:
         :meth:`snowflake.snowpark.Session.query_history`.
     """
 
-    def __init__(self, session: "snowflake.snowpark.session.Session") -> None:
+    def __init__(
+        self,
+        session: "snowflake.snowpark.session.Session",
+        include_describe: bool = False,
+    ) -> None:
         self.session = session
         self._queries: List[QueryRecord] = []
+        self._include_describe = include_describe
 
     def __enter__(self):
         return self
@@ -37,3 +43,7 @@ def _add_query(self, query_record: QueryRecord):
     @property
     def queries(self) -> List[QueryRecord]:
         return self._queries
+
+    @property
+    def include_describe(self) -> bool:
+        return self._include_describe

src/snowflake/snowpark/session.py

@@ -3456,16 +3456,21 @@ 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:
+        Args:
+            include_describe: Include query notifications for describe queries
+
+        >>> with session.query_history(True) as query_history:
         ...     df = session.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"])
         ...     df = df.filter(df.a == 1)
         ...     res = df.collect()
-        >>> assert len(query_history.queries) == 1
+        >>> assert len(query_history.queries) == 2
+        >>> assert query_history.queries[0].is_describe
+        >>> assert not query_history.queries[1].is_describe
         """
-        query_listener = QueryHistory(self)
+        query_listener = QueryHistory(self, include_describe)
         self._conn.add_query_listener(query_listener)
         return query_listener
 

tests/integ/modin/sql_counter.py

@@ -14,7 +14,7 @@
 from decorator import decorator
 from pandas._typing import Scalar
 
-from snowflake.snowpark import QueryRecord
+from snowflake.snowpark.query_history import QueryRecord
 from snowflake.snowpark.session import Session
 from tests.utils import IS_IN_STORED_PROC
 
@@ -43,6 +43,7 @@
 UDTF_COUNT_PARAMETER = "udtf_count"
 SELECT_COUNT_PARAMETER = "select_count"
 UNION_COUNT_PARAMETER = "union_count"
+DESCRIBE_COUNT_PARAMETER = "describe_count"
 EXPECT_HIGH_COUNT = "expect_high_count"
 HIGH_COUNT_REASON = "high_count_reason"
 
@@ -54,6 +55,7 @@
     UDTF_COUNT_PARAMETER,
     SELECT_COUNT_PARAMETER,
     UNION_COUNT_PARAMETER,
+    DESCRIBE_COUNT_PARAMETER,
 ]
 BOOL_PARAMETERS = [EXPECT_HIGH_COUNT]
 
@@ -156,6 +158,11 @@ def __init__(
             # Add SqlCounter as a snowpark query listener.
             self.session._conn.add_query_listener(self)
 
+    # The query history listener will include describe queries if this is true.
+    @property
+    def include_describe(self) -> bool:
+        return True
+
     @staticmethod
     def set_record_mode(record_mode):
         """Record mode means the SqlCounter does not assert any results, but rather collects them so they can
@@ -251,7 +258,7 @@ def expects(self, **kwargs):
 
         # If there are any failures, print out all the captured queries so clear which are being counted.
         if failed:
-            title = f"{'='*20} SqlCounter Captured Queries {'='*20}"
+            title = f"\n{'='*20} SqlCounter Captured Queries {'='*20}\n"
             print(title, file=sys.stderr)
             for query in self._get_actual_queries():
                 print(query, file=sys.stderr)
@@ -268,7 +275,12 @@ def _get_actual_queries(self):
                         for fw in FILTER_OUT_QUERIES
                     ]
                 ),
-                list(map(lambda q: q.sql_text, self._queries)),
+                list(
+                    map(
+                        lambda q: q.sql_text,
+                        [q for q in self._queries if not q.is_describe],
+                    )
+                ),
             )
         )
 
@@ -329,6 +341,9 @@ def actual_select_count(self):
     def actual_union_count(self):
         return self._count_instances_by_query_substr(contains=[UNION])
 
+    def actual_describe_count(self):
+        return len([q for q in self._queries if q.is_describe])
+
     def get_actual_counts(self):
         """Retrieve all actual counts so far."""
         actual_counts = {}

tests/integ/modin/test_sql_counter.py

@@ -80,22 +80,22 @@ def test_sql_counter_with_context_manager_inside_loop():
 
 @sql_count_checker(no_check=True)
 def test_sql_counter_with_multiple_checks():
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, describe_count=1):
         df = pd.DataFrame({"a": [1, 2, 3]})
         assert len(df) == 3
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, describe_count=1):
         df = pd.DataFrame({"b": [4, 5, 6]})
         assert len(df) == 3
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, describe_count=1):
         df = pd.DataFrame({"c": [7, 8, 9]})
         assert len(df) == 3
 
 
 @sql_count_checker(no_check=True)
 def test_sql_counter_with_context_manager_outside_loop():
-    sc = SqlCounter(query_count=3)
+    sc = SqlCounter(query_count=3, describe_count=3)
     sc.__enter__()
     for _ in range(3):
         df = pd.DataFrame({"a": [1, 2, 3]})

tests/integ/test_query_history.py

@@ -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):
@@ -57,6 +71,16 @@ def test_query_history_two_listeners(session):
 
 
 def test_query_history_multiple_actions(session):
+    with session.query_history(True) as query_history:
+        df = session.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"])
+        df = df.filter(df.a == 1)
+        df.collect()
+
+    assert len(query_history.queries) == 3
+    assert query_history.queries[0].is_describe
+    assert query_history.queries[1].is_describe
+    assert not query_history.queries[2].is_describe
+
     with session.query_history() as query_listener:
         session.sql("select 0").collect()
         session.sql("select 1").collect()