Query: Adds Query Advisor SDK capabilities

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 4.15.1 (Unreleased)
 
 #### Features Added
+* Added support for Query Advisor feature - See [PR 45331](https://github.com/Azure/azure-sdk-for-python/pull/45331)
 * Added `get_response_headers()` and `get_last_response_headers()` methods to the `CosmosItemPaged` and `CosmosAsyncItemPaged` objects returned by `query_items()`, allowing access to response headers from query operations. See [PR 44593](https://github.com/Azure/azure-sdk-for-python/pull/44593)
 
 #### Breaking Changes

sdk/cosmos/azure-cosmos/MANIFEST.in

@@ -4,4 +4,5 @@ include azure/__init__.py
 recursive-include samples *.py *.md
 recursive-include tests *.py
 include azure/cosmos/py.typed
+recursive-include azure/cosmos/_query_advisor *.json
 recursive-include doc *.rst

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -359,6 +359,9 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("populateIndexMetrics"):
         headers[http_constants.HttpHeaders.PopulateIndexMetrics] = options["populateIndexMetrics"]
 
+    if options.get("populateQueryAdvice"):
+        headers[http_constants.HttpHeaders.PopulateQueryAdvice] = options["populateQueryAdvice"]
+
     if options.get("responseContinuationTokenLimitInKb"):
         headers[http_constants.HttpHeaders.ResponseContinuationTokenLimitInKb] = options[
             "responseContinuationTokenLimitInKb"]

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -71,6 +71,7 @@
 from ._request_object import RequestObject
 from ._retry_utility import ConnectionRetryPolicy
 from ._routing import routing_map_provider, routing_range
+from ._query_advisor import get_query_advice_info
 from ._inference_service import _InferenceService
 from .documents import ConnectionPolicy, DatabaseAccount
 from .partition_key import (
@@ -3379,6 +3380,14 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                     response_hook(last_response_headers, partial_result)
             # if the prefix partition query has results lets return it
             if results:
+                if last_response_headers.get(http_constants.HttpHeaders.IndexUtilization) is not None:
+                    index_metrics_raw = last_response_headers[http_constants.HttpHeaders.IndexUtilization]
+                    last_response_headers[http_constants.HttpHeaders.IndexUtilization] = (
+                         _utils.get_index_metrics_info(index_metrics_raw))
+                if last_response_headers.get(http_constants.HttpHeaders.QueryAdvice) is not None:
+                    query_advice_raw = last_response_headers[http_constants.HttpHeaders.QueryAdvice]
+                    last_response_headers[http_constants.HttpHeaders.QueryAdvice] = (
+                        get_query_advice_info(query_advice_raw))
                 return __GetBodiesFromQueryResult(results), last_response_headers
 
         result, last_response_headers = self.__Post(path, request_params, query, req_headers, **kwargs)
@@ -3388,6 +3397,9 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
             INDEX_METRICS_HEADER = http_constants.HttpHeaders.IndexUtilization
             index_metrics_raw = last_response_headers[INDEX_METRICS_HEADER]
             last_response_headers[INDEX_METRICS_HEADER] = _utils.get_index_metrics_info(index_metrics_raw)
+        if last_response_headers.get(http_constants.HttpHeaders.QueryAdvice) is not None:
+            query_advice_raw = last_response_headers[http_constants.HttpHeaders.QueryAdvice]
+            last_response_headers[http_constants.HttpHeaders.QueryAdvice] = get_query_advice_info(query_advice_raw)
         if response_headers_list is not None:
             response_headers_list.append(last_response_headers.copy())
         if response_hook:

sdk/cosmos/azure-cosmos/azure/cosmos/_query_advisor/__init__.py

@@ -0,0 +1,18 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+"""Query Advisor module for processing query optimization advice from Azure Cosmos DB."""
+
+from ._query_advice import QueryAdvice, QueryAdviceEntry
+from ._rule_directory import RuleDirectory
+from ._get_query_advice_info import get_query_advice_info
+
+__all__ = [
+    "QueryAdvice",
+    "QueryAdviceEntry",
+    "RuleDirectory",
+    "get_query_advice_info",
+]

sdk/cosmos/azure-cosmos/azure/cosmos/_query_advisor/_get_query_advice_info.py

@@ -0,0 +1,34 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+"""Function for processing query advice response headers."""
+
+from typing import Optional
+
+from ._query_advice import QueryAdvice
+
+def get_query_advice_info(header_value: Optional[str]) -> str:
+    """Process a query advice response header into a formatted human-readable string.
+
+    Takes the raw ``x-ms-cosmos-query-advice`` response header (URL-encoded JSON),
+    decodes it, parses the query advice entries, enriches them with human-readable
+    messages from the rule directory, and returns a formatted multi-line string.
+
+    :param str header_value: The raw query advice response header value (URL-encoded JSON).
+    :returns: Formatted string with query advice entries, or empty string if parsing fails.
+    :rtype: str
+    """
+    if header_value is None:
+        return ""
+
+    # Parse the query advice from the header
+    query_advice = QueryAdvice.try_create_from_string(header_value)
+
+    if query_advice is None:
+        return ""
+
+    # Format as string
+    return str(query_advice)

sdk/cosmos/azure-cosmos/azure/cosmos/_query_advisor/_query_advice.py

@@ -0,0 +1,151 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+"""Query advice classes for parsing and formatting query optimization recommendations."""
+
+import json
+import logging
+from typing import Any, Dict, List, Optional
+from urllib.parse import unquote
+
+from ._rule_directory import RuleDirectory
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class QueryAdviceEntry:
+    """Represents a single query advice entry.
+
+    Each entry contains a rule ID and optional parameters that provide
+    specific guidance for query optimization.
+    """
+
+    def __init__(self, rule_id: str, parameters: Optional[List[str]] = None) -> None:
+        """Initialize a query advice entry.
+
+        :param str rule_id: The rule identifier (e.g., ``QA1000``).
+        :param parameters: Optional list of parameters for the rule message.
+        :type parameters: list[str] or None
+        """
+        self.id = rule_id
+        self.parameters = parameters or []
+
+    def __str__(self) -> str:
+        """Format the query advice entry as a human-readable string.
+
+        :returns: Formatted string with rule ID, message, and documentation link,
+            or empty string if the rule identifier is missing.
+        :rtype: str
+        """
+        if self.id is None:
+            return ""
+
+        rule_directory = RuleDirectory()
+        message = rule_directory.get_rule_message(self.id)
+        if message is None:
+            # Unknown rule — log it and return the public doc link as the fallback.
+            fallback_url = f"{rule_directory.url_prefix}{self.id}"
+            _LOGGER.warning(
+                "Unknown Query Advisor rule '%s'. For more information, please visit %s",
+                self.id,
+                fallback_url,
+            )
+            return f"{self.id}: For more information, please visit {fallback_url}"
+
+        # Format: {id}: {message} For more information, please visit {url_prefix}{id}
+        result = f"{self.id}: "
+
+        # Format message with parameters if available
+        if self.parameters:
+            try:
+                result += message.format(*self.parameters)
+            except (IndexError, KeyError):
+                # If formatting fails, use message as-is
+                result += message
+        else:
+            result += message
+
+        # Add documentation link
+        result += f" For more information, please visit {rule_directory.url_prefix}{self.id}"
+
+        return result
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "QueryAdviceEntry":
+        """Create a QueryAdviceEntry from a dictionary.
+
+        :param data: Dictionary with "Id" and optional "Params" keys.
+        :type data: dict[str, any]
+        :returns: QueryAdviceEntry instance.
+        :rtype: ~azure.cosmos._query_advisor._query_advice.QueryAdviceEntry
+        """
+        rule_id = data.get("Id", "")
+        parameters = data.get("Params", [])
+        return cls(rule_id, parameters)
+
+
+class QueryAdvice:
+    """Collection of query advice entries.
+
+    Represents the complete query advice response from Azure Cosmos DB,
+    containing one or more optimization recommendations.
+    """
+
+    def __init__(self, entries: Optional[List[QueryAdviceEntry]] = None) -> None:
+        """Initialize query advice with a list of entries.
+
+        :param entries: List of QueryAdviceEntry objects.
+        :type entries: list[~azure.cosmos._query_advisor._query_advice.QueryAdviceEntry] or None
+        """
+        self.entries = [e for e in (entries or []) if e is not None]
+
+    def __str__(self) -> str:
+        """Format all query advice entries as a multi-line string.
+
+        :returns: Formatted string with each entry on a separate line.
+        :rtype: str
+        """
+        if not self.entries:
+            return ""
+
+        lines = []
+
+        for entry in self.entries:
+            formatted = str(entry)
+            if formatted:
+                lines.append(formatted)
+
+        return "\n".join(lines)
+
+    @classmethod
+    def try_create_from_string(cls, response_header: Optional[str]) -> Optional["QueryAdvice"]:
+        """Parse query advice from a URL-encoded JSON response header.
+
+        :param response_header: URL-encoded JSON string from the response header.
+        :type response_header: str or None
+        :returns: QueryAdvice instance if parsing succeeds, None otherwise.
+        :rtype: ~azure.cosmos._query_advisor._query_advice.QueryAdvice or None
+        """
+        if response_header is None:
+            return None
+
+        try:
+            # URL-decode the header value
+            decoded_string = unquote(response_header)
+
+            # Parse JSON into list of entry dictionaries
+            data = json.loads(decoded_string)
+
+            if not isinstance(data, list):
+                return None
+
+            # Convert dictionaries to QueryAdviceEntry objects
+            entries = [QueryAdviceEntry.from_dict(item) for item in data if isinstance(item, dict)]
+
+            return cls(entries)
+        except (json.JSONDecodeError, ValueError, AttributeError) as e:
+            _LOGGER.warning("Failed to parse query advice from response header: %s", e)
+            return None

sdk/cosmos/azure-cosmos/azure/cosmos/_query_advisor/_rule_directory.py

@@ -0,0 +1,81 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+"""Rule directory singleton for loading and accessing query advice rules."""
+
+import json
+import logging
+from importlib.resources import files
+from typing import Any, Dict, Optional
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class RuleDirectory:
+    """Singleton for loading and accessing query advice rules.
+
+    The rule directory lazy-loads the query_advice_rules.json file
+    and provides access to rule messages and URL prefix.
+    Uses importlib.resources so it works correctly in all packaging
+    scenarios including zip-safe wheels.
+    """
+
+    _instance: Optional["RuleDirectory"] = None
+
+    def __new__(cls) -> "RuleDirectory":
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self) -> None:
+        # Guard so the singleton body only runs once.
+        if getattr(self, "_initialized", False):
+            return
+
+        self._initialized: bool = True
+        self._rules: Dict[str, Dict[str, Any]] = {}
+        self._url_prefix: str = ""
+        self._load_rules()
+
+    def _load_rules(self) -> None:
+        """Load rules from the bundled JSON resource."""
+        try:
+            resource_text = (
+                files(__package__)
+                .joinpath("query_advice_rules.json")
+                .read_text(encoding="utf-8")
+            )
+            data = json.loads(resource_text)
+            self._url_prefix = data.get("url_prefix", "")
+            self._rules = data.get("rules", {})
+        except Exception:  # pylint: disable=broad-except
+            # Fall back to empty rules so query execution
+            # is never blocked by an inability to load advice text.
+            _LOGGER.warning("Failed to load query_advice_rules.json, falling back to empty rules", exc_info=True)
+            self._url_prefix = (
+                "https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/query/queryadvisor/"
+            )
+            self._rules = {}
+
+    @property
+    def url_prefix(self) -> str:
+        """Get the URL prefix for documentation links.
+
+        :rtype: str
+        """
+        return self._url_prefix
+
+    def get_rule_message(self, rule_id: str) -> Optional[str]:
+        """Get the message for a given rule ID.
+
+        :param str rule_id: The rule identifier (e.g., ``QA1000``).
+        :returns: The rule message, or ``None`` if the rule is not found.
+        :rtype: str or None
+        """
+        rule = self._rules.get(rule_id)
+        if rule:
+            return rule.get("message")
+        return None