fix(chronicle): prevent response data loss in nl_search and search_udm when :udmSearch returns stats

README.md

@@ -1265,6 +1265,8 @@ The natural language search feature supports various query patterns:
 - "Show login failures in the last hour"
 - "Display connections to IP address 192.168.1.100"
 
+When a natural language prompt translates into a statistical aggregation query, the returned dictionary includes a `stats` key with `columns`, `rows`, and `total_rows` matching `get_stats()`, while `events` is empty.
+
 If the natural language cannot be translated to a valid UDM query, an `APIError` will be raised with a message indicating that no valid query could be generated.
 
 ### Entity Summary

src/secops/chronicle/client.py

@@ -969,13 +969,18 @@ def search_udm(
             max_attempts: Maximum number of polling attempts (default: 30)
             timeout: Timeout in seconds for each API request (default: 30)
             debug: Print debug information during execution
-            as_list: If True, return a list of events instead of a dict
-                with events list and nextPageToken.
+            as_list: If True, return only the list of events; statistical
+                results in stats are omitted when as_list is True.
 
         Returns:
-            If as_list is True: List of Events.
-            If as_list is False: Dict with event list, total number of event and
-                flag to check if more data is available.
+            If as_list is True: List of matched UDM events, or an empty list
+                if the query returned only statistical results.
+            If as_list is False: Dict with events, total_events,
+                more_data_available, and an optional stats dict containing
+                columns, rows, and total_rows matching get_stats() when the
+                query returns statistical results. For dedicated statistical
+                queries with custom max_values or a 120s default timeout,
+                use get_stats().
 
         Raises:
             APIError: If the API request fails
@@ -3360,7 +3365,9 @@ def nl_search(
             max_attempts: Maximum number of polling attempts
 
         Returns:
-            Dict containing the search results with events
+            Dict with events, total_events, more_data_available, and an
+            optional stats dict containing columns, rows, and total_rows
+            when the translated query returns statistical results.
 
         Raises:
             APIError: If the API request fails

src/secops/chronicle/nl_search.py

@@ -72,7 +72,9 @@ def nl_search(
         max_attempts: Maximum number of polling attempts
 
     Returns:
-        Dict containing the search results with events
+        Dict with events, total_events, more_data_available, and an
+        optional stats dict containing columns, rows, and total_rows
+        when the translated query returns statistical results.
 
     Raises:
         APIError: If the API request fails after retries

src/secops/chronicle/search.py

@@ -18,6 +18,7 @@
 from typing import Any, TYPE_CHECKING
 
 from secops.chronicle.models import APIVersion
+from secops.chronicle.stats import process_stats_results
 from secops.chronicle.utils.request_utils import (
     chronicle_request,
 )
@@ -51,12 +52,18 @@ def search_udm(
                         for backwards compatibility)
         timeout: Timeout in seconds for each API request (default: 30)
         debug: Print debug information during execution
-        as_list: Whether to return results as a list or dictionary
+        as_list: If True, return only the list of events; statistical
+            results in stats are omitted when as_list is True.
 
     Returns:
-        If as_list is True: List of Events.
-        If as_list is False: Dict with event list, total number of event and
-            flag to check if more data is available.
+        If as_list is True: List of matched UDM events, or an empty list
+            if the query returned only statistical results.
+        If as_list is False: Dict with events, total_events,
+            more_data_available, and an optional stats dict containing
+            columns, rows, and total_rows matching get_stats() when the
+            query returns statistical results. For dedicated statistical
+            queries with custom max_values or a 120s default timeout,
+            use get_stats().
 
     Raises:
         APIError: If the API request fails
@@ -93,8 +100,17 @@ def search_udm(
         return result.get("events", [])
 
     events = result.get("events", [])
-    return {
+    raw_stats = result.get("stats")
+    has_stats_dict = isinstance(raw_stats, dict)
+    more_data_available = bool(
+        result.get("moreDataAvailable", False)
+        or (has_stats_dict and raw_stats.get("tooManyResults", False))
+    )
+    response: dict[str, Any] = {
         "events": events,
         "total_events": len(events),
-        "more_data_available": result.get("moreDataAvailable", False),
+        "more_data_available": more_data_available,
     }
+    if has_stats_dict:
+        response["stats"] = process_stats_results(raw_stats)
+    return response

tests/chronicle/test_nl_search.py

@@ -18,6 +18,7 @@
 from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock, patch
 from secops.chronicle.nl_search import translate_nl_to_udm, nl_search
+from secops.chronicle.search import search_udm
 from secops.exceptions import APIError
 import unittest.mock
 
@@ -244,3 +245,59 @@ def test_nl_search_max_retries_exceeded(
 
     # Verify sleep was called 10 times
     assert mock_sleep.call_count == 10
+
+
+def test_nl_search_preserves_stats_from_outcome_query(mock_client):
+    """Test nl_search returns stats when translation emits an outcome: query."""
+    mock_client.search_udm.side_effect = lambda **kwargs: search_udm(
+        mock_client, **kwargs
+    )
+
+    translate_resp = MagicMock()
+    translate_resp.status_code = 200
+    translate_resp.json.return_value = {
+        "query": (
+            'metadata.event_type = "USER_LOGIN"\n'
+            "match:\n  target.user.userid\n"
+            "outcome:\n  $login_count = count(metadata.id)"
+        )
+    }
+
+    udm_search_resp = MagicMock()
+    udm_search_resp.status_code = 200
+    udm_search_resp.json.return_value = {
+        "stats": {
+            "results": [
+                {
+                    "column": "target.user.userid",
+                    "values": [{"value": {"stringVal": "alice"}}],
+                },
+                {
+                    "column": "login_count",
+                    "values": [{"value": {"int64Val": "15"}}],
+                },
+            ],
+            "tooManyResults": False,
+        }
+    }
+
+    mock_client.session.request.side_effect = [translate_resp, udm_search_resp]
+
+    start_time = datetime.now(timezone.utc) - timedelta(hours=24)
+    end_time = datetime.now(timezone.utc)
+
+    result = nl_search(
+        mock_client,
+        "count user logins by user in the last 24 hours",
+        start_time,
+        end_time,
+    )
+
+    assert result["events"] == []
+    assert result["total_events"] == 0
+    assert result["more_data_available"] is False
+    assert result["stats"] == {
+        "columns": ["target.user.userid", "login_count"],
+        "rows": [{"target.user.userid": "alice", "login_count": 15}],
+        "total_rows": 1,
+    }

tests/chronicle/test_search.py

@@ -77,6 +77,7 @@ def test_search_udm_defaults_when_keys_missing(
         self.assertEqual(result["events"], [])
         self.assertEqual(result["total_events"], 0)
         self.assertFalse(result["more_data_available"])
+        self.assertNotIn("stats", result)
 
     @mock.patch("secops.chronicle.search.chronicle_request")
     @mock.patch("builtins.print")
@@ -194,6 +195,164 @@ def test_search_udm_empty_events_list(
         self.assertEqual(result["events"], [])
         self.assertEqual(result["total_events"], 0)
         self.assertFalse(result["more_data_available"])
+        self.assertNotIn("stats", result)
+
+    @mock.patch("secops.chronicle.search.chronicle_request")
+    def test_search_udm_includes_processed_stats_when_present(
+        self, mock_chronicle_request: mock.MagicMock
+    ) -> None:
+        """Test that stats payloads from outcome: queries are preserved."""
+        mock_chronicle_request.return_value = {
+            "stats": {
+                "results": [
+                    {
+                        "column": "principal.ip",
+                        "values": [
+                            {"value": {"stringVal": "10.0.0.1"}},
+                            {"value": {"stringVal": "10.0.0.2"}},
+                        ],
+                    },
+                    {
+                        "column": "login_count",
+                        "values": [
+                            {"value": {"int64Val": "42"}},
+                            {"value": {"int64Val": "7"}},
+                        ],
+                    },
+                ],
+                "tooManyResults": True,
+            }
+        }
+
+        result = search_udm(
+            client=self.client,
+            query=(
+                'metadata.event_type = "USER_LOGIN"\n'
+                "match:\n  principal.ip\n"
+                "outcome:\n  $login_count = count(metadata.id)"
+            ),
+            start_time=self.start_time,
+            end_time=self.end_time,
+        )
+
+        self.assertEqual(result["events"], [])
+        self.assertEqual(result["total_events"], 0)
+        self.assertTrue(result["more_data_available"])
+        self.assertIn("stats", result)
+        self.assertEqual(
+            result["stats"],
+            {
+                "columns": ["principal.ip", "login_count"],
+                "rows": [
+                    {"principal.ip": "10.0.0.1", "login_count": 42},
+                    {"principal.ip": "10.0.0.2", "login_count": 7},
+                ],
+                "total_rows": 2,
+            },
+        )
+
+    @mock.patch("secops.chronicle.search.chronicle_request")
+    def test_search_udm_empty_stats_payload_preserved(
+        self, mock_chronicle_request: mock.MagicMock
+    ) -> None:
+        """Test zero-row outcome: stats payloads remain distinguishable."""
+        mock_chronicle_request.return_value = {"stats": {}}
+
+        empty_result = search_udm(
+            client=self.client,
+            query=(
+                'metadata.event_type = "USER_LOGIN"\n'
+                "outcome:\n  $event_count = count(metadata.id)"
+            ),
+            start_time=self.start_time,
+            end_time=self.end_time,
+        )
+
+        self.assertFalse(empty_result["more_data_available"])
+        self.assertEqual(
+            empty_result["stats"],
+            {"columns": [], "rows": [], "total_rows": 0},
+        )
+
+        mock_chronicle_request.return_value = {
+            "stats": {
+                "results": [
+                    {"column": "principal.ip"},
+                    {"column": "login_count"},
+                ]
+            }
+        }
+
+        header_only_result = search_udm(
+            client=self.client,
+            query=(
+                'metadata.event_type = "USER_LOGIN"\n'
+                "match:\n  principal.ip\n"
+                "outcome:\n  $login_count = count(metadata.id)"
+            ),
+            start_time=self.start_time,
+            end_time=self.end_time,
+        )
+
+        self.assertFalse(header_only_result["more_data_available"])
+        self.assertEqual(
+            header_only_result["stats"],
+            {
+                "columns": ["principal.ip", "login_count"],
+                "rows": [],
+                "total_rows": 0,
+            },
+        )
+
+    @mock.patch("secops.chronicle.search.chronicle_request")
+    def test_search_udm_ignores_non_dict_stats_value(
+        self, mock_chronicle_request: mock.MagicMock
+    ) -> None:
+        """Test malformed non-dict stats payloads do not raise errors."""
+        mock_chronicle_request.return_value = {
+            "events": [],
+            "stats": "unexpected",
+        }
+
+        result = search_udm(
+            client=self.client,
+            query="test",
+            start_time=self.start_time,
+            end_time=self.end_time,
+        )
+
+        self.assertEqual(result["events"], [])
+        self.assertFalse(result["more_data_available"])
+        self.assertNotIn("stats", result)
+
+    @mock.patch("secops.chronicle.search.chronicle_request")
+    def test_search_udm_as_list_with_stats_returns_empty_events_list(
+        self, mock_chronicle_request: mock.MagicMock
+    ) -> None:
+        """Test as_list=True returns empty event list on stats responses."""
+        mock_chronicle_request.return_value = {
+            "stats": {
+                "results": [
+                    {
+                        "column": "event_count",
+                        "values": [{"value": {"int64Val": "5"}}],
+                    }
+                ]
+            }
+        }
+
+        result = search_udm(
+            client=self.client,
+            query=(
+                'metadata.event_type = "USER_LOGIN"\n'
+                "outcome:\n  $event_count = count(metadata.id)"
+            ),
+            start_time=self.start_time,
+            end_time=self.end_time,
+            as_list=True,
+        )
+
+        self.assertEqual(result, [])
 
 
 if __name__ == "__main__":