Merge branch 'm-kovalsky/enablecaching'

Author: m-kovalsky

src/sempy_labs/semantic_model/__init__.py

@@ -3,10 +3,14 @@
     set_endorsement,
     make_discoverable,
 )
+from ._caching import (
+    enable_query_caching,
+)
 
 
 __all__ = [
     "approved_for_copilot",
     "set_endorsement",
     "make_discoverable",
+    "enable_query_caching",
 ]

src/sempy_labs/semantic_model/_caching.py

@@ -0,0 +1,60 @@
+import requests
+from uuid import UUID
+from typing import Optional
+from sempy_labs._helper_functions import (
+    get_pbi_token_headers,
+    _get_url_prefix,
+    get_model_id,
+    resolve_item_name_and_id,
+    resolve_workspace_name_and_id,
+)
+from sempy.fabric.exceptions import FabricHTTPException
+from sempy._utils._log import log
+import sempy_labs._icons as icons
+
+
+@log
+def enable_query_caching(
+    dataset: str | UUID, workspace: Optional[str | UUID] = None, enable: bool = True,
+):
+    """
+    Enables or disables `query caching <http://aka.ms/queryCaching>`_ for a semantic model.
+
+    Parameters
+    ----------
+    dataset : str | uuid.UUID
+        Name or ID of the semantic model.
+    workspace : str | uuid.UUID, default=None
+        The workspace name or ID.
+        Defaults to None which resolves to the workspace of the attached lakehouse
+        or if no lakehouse attached, resolves to the workspace of the notebook.
+    enable : bool, default=True
+        Set to True to enable query caching, or False to disable it.
+    """
+    prefix = _get_url_prefix()
+    headers = get_pbi_token_headers()
+    (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
+    (item_name, item_id) = resolve_item_name_and_id(
+        item=dataset, type="SemanticModel", workspace=workspace_id
+    )
+    model_id = get_model_id(item_id=item_id, headers=headers, prefix=prefix)
+    if model_id is None:
+        raise ValueError(f"Failed to retrieve model ID for semantic model '{item_name}'")
+
+    caching_map = {
+        True: 2,
+        False: 1,
+    }
+
+    payload = {"queryCachingState": caching_map.get(enable)}
+
+    response = requests.post(
+        f"{prefix}/metadata/models/{model_id}/caching", headers=headers, json=payload
+    )
+
+    if response.status_code != 204:
+        raise FabricHTTPException(response)
+
+    print(
+        f"{icons.green_dot} Query caching has been {'enabled' if enable else 'disabled'} for the '{item_name}' semantic model within the '{workspace_name}' workspace."
+    )

src/sempy_labs/tom/_model.py

@@ -814,7 +814,7 @@ def delete_user_defined_function(self, name: str):
 
     def set_user_defined_function(self, name: str, expression: str):
         """
-        Sets the definition of a `user-defined <https://learn.microsoft.com/en-us/dax/best-practices/dax-user-defined-functions#using-model-explorer>_` function within the semantic model. This function requires that the compatibility level is at least 1702.
+        Sets the definition of a `user-defined <https://learn.microsoft.com/en-us/dax/best-practices/dax-user-defined-functions#using-model-explorer>`_ function within the semantic model. This function requires that the compatibility level is at least 1702.
 
         Parameters
         ----------