Query: Adds Query Advisor SDK capabilities

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 4.15.1 (Unreleased)
 
 #### Features Added
+* Added Query Advisor support for Python SDK. See [PR 45331](https://github.com/Azure/azure-sdk-for-python/pull/45331)
 
 #### 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 (
@@ -3376,6 +3377,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_hook:
             response_hook(last_response_headers, result)
 

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,35 @@
+# -------------------------------------------------------------------------
+# 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
+
+# cspell:ignore STARTSWTIH
+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 query_advice.to_string()

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

@@ -0,0 +1,142 @@
+# -------------------------------------------------------------------------
+# 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
+from typing import Any, Dict, List, Optional
+from urllib.parse import unquote
+
+from ._rule_directory import RuleDirectory
+
+
+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.
+
+        Args:
+            rule_id: The rule identifier (e.g., "QA1000")
+            parameters: Optional list of parameters for the rule message
+        """
+        self.id = rule_id
+        self.parameters = parameters or []
+
+    def to_string(self, rule_directory: RuleDirectory) -> Optional[str]:
+        """Format the query advice entry as a human-readable string.
+
+        :param rule_directory: Rule directory instance for looking up messages.
+        :type rule_directory: ~azure.cosmos._query_advisor._rule_directory.RuleDirectory
+        :returns: Formatted string with rule ID, message, and documentation link,
+            or None if the rule message cannot be found.
+        :rtype: str or None
+        """
+        if self.id is None:
+            return None
+
+        message = rule_directory.get_rule_message(self.id)
+        if message is None:
+            return None
+
+        # 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.
+
+        Args:
+            entries: List of QueryAdviceEntry objects
+        """
+        self.entries = [e for e in (entries or []) if e is not None]
+
+    def to_string(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 ""
+
+        rule_directory = RuleDirectory()
+        lines = []
+
+        for entry in self.entries:
+            formatted = entry.to_string(rule_directory)
+            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):
+            return None

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

@@ -0,0 +1,77 @@
+# -------------------------------------------------------------------------
+# 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
+from importlib.resources import files
+from typing import Any, Dict, Optional
+
+
+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
+            # Silently fall back to empty rules so query execution
+            # is never blocked by an inability to load advice text.
+            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

sdk/cosmos/azure-cosmos/azure/cosmos/_query_advisor/query_advice_rules.json

@@ -0,0 +1,55 @@
+{
+  "url_prefix": "https://aka.ms/CosmosDB/QueryAdvisor/",
+  "rules": {
+    "QA1000": {
+      "name": "PartialArrayContains",
+      "description": "Query uses ARRAY_CONTAINS with the third argument set to true.",
+      "message": "Instead of ARRAY_CONTAINS, consider using EXISTS with a subquery, which may improve performance."
+    },
+    "QA1001": {
+      "name": "DistinctAndJoin",
+      "description": "Query uses Distinct and Join.",
+      "message": "Instead of DISTINCT with a JOIN, consider using EXISTS with a subquery, which may improve performance."
+    },
+    "QA1002": {
+      "name": "Contains",
+      "description": "Query uses CONTAINS.",
+      "message": "If you are matching on a string prefix, consider using STARTSWITH."
+    },
+    "QA1003": {
+      "name": "CaseInsensitiveStartsWithOrStringEquals",
+      "description": "Query uses case-insensitive string search functions STARTSWITH or StringEquals.",
+      "message": "Instead of case-insensitive string search, consider creating a computed property with LOWER on the string field, which may improve performance."
+    },
+    "QA1004": {
+      "name": "CaseInsensitiveEndsWith",
+      "description": "Query uses case-insensitive ENDSWITH.",
+      "message": "Instead of case-insensitive ENDSWITH, consider creating a computed property with REVERSE on the string field, and use STARTSWITH for comparison, which may improve performance."
+    },
+    "QA1005": {
+      "name": "GroupByComputedProperty",
+      "description": "Query uses deterministic scalar expressions in Group By clause.",
+      "message": "Instead of using scalar expressions in GROUP BY clause, consider creating computed properties of these expressions, which may improve performance."
+    },
+    "QA1006": {
+      "name": "UpperLowerComparison",
+      "description": "Query uses Upper or Lower string comparison.",
+      "message": "Consider defining a computed property on the UPPER/LOWER function expression."
+    },
+    "QA1007": {
+      "name": "GetCurrentDateTime",
+      "description": "Query uses GetCurrentDateTime.",
+      "message": "Consider using GetCurrentDateTimeStatic instead of GetCurrentDateTime in the WHERE clause."
+    },
+    "QA1008": {
+      "name": "GetCurrentTicks",
+      "description": "Query uses GetCurrentTicks.",
+      "message": "Consider using GetCurrentTicksStatic instead of GetCurrentTicks in the WHERE clause."
+    },
+    "QA1009": {
+      "name": "GetCurrentTimestamp",
+      "description": "Query uses GetCurrentTimestamp.",
+      "message": "Consider using GetCurrentTimestampStatic instead of GetCurrentTimestamp in the WHERE clause."
+    }
+  }
+}