semantic-link-labs 0.9.1__py3-none-any.whl → 0.9.3__py3-none-any.whl

sempy_labs/_dataflows.py

@@ -3,10 +3,12 @@ import pandas as pd
 from sempy_labs._helper_functions import (
     resolve_workspace_name_and_id,
     _is_valid_uuid,
+    _update_dataframe_datatypes,
+    _base_api,
+    _create_dataframe,
 )
 from typing import Optional, Tuple
 import sempy_labs._icons as icons
-from sempy.fabric.exceptions import FabricHTTPException
 from uuid import UUID
 
 
@@ -28,14 +30,17 @@ def list_dataflows(workspace: Optional[str | UUID] = None):
     """
 
     (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
-    client = fabric.PowerBIRestClient()
-    response = client.get(f"/v1.0/myorg/groups/{workspace_id}/dataflows")
-    if response.status_code != 200:
-        raise FabricHTTPException(response)
 
-    df = pd.DataFrame(
-        columns=["Dataflow Id", "Dataflow Name", "Configured By", "Users", "Generation"]
-    )
+    columns = {
+        "Dataflow Id": "string",
+        "Dataflow Name": "string",
+        "Configured By": "string",
+        "Users": "string",
+        "Generation": "int",
+    }
+    df = _create_dataframe(columns=columns)
+
+    response = _base_api(request=f"/v1.0/myorg/groups/{workspace_id}/dataflows")
 
     data = []  # Collect rows here
 
@@ -51,7 +56,8 @@ def list_dataflows(workspace: Optional[str | UUID] = None):
 
     if data:
         df = pd.DataFrame(data)
-        df["Generation"] = df["Generation"].astype(int)
+
+        _update_dataframe_datatypes(dataframe=df, column_map=columns)
 
     return df
 
@@ -85,16 +91,14 @@ def assign_workspace_to_dataflow_storage(
         )
 
     dataflow_storage_id = df_filt["Dataflow Storage Account ID"].iloc[0]
-    client = fabric.PowerBIRestClient()
-
-    request_body = {"dataflowStorageId": dataflow_storage_id}
+    payload = {"dataflowStorageId": dataflow_storage_id}
 
-    response = client.post(
-        f"/v1.0/myorg/groups/{workspace_id}/AssignToDataflowStorage", json=request_body
+    _base_api(
+        request=f"/v1.0/myorg/groups/{workspace_id}/AssignToDataflowStorage",
+        method="post",
+        payload=payload,
     )
 
-    if response.status_code != 200:
-        raise FabricHTTPException(response)
     print(
         f"{icons.green_dot} The '{dataflow_storage_account}' dataflow storage account has been assigned to the '{workspace_name}' workspacce."
     )
@@ -112,17 +116,14 @@ def list_dataflow_storage_accounts() -> pd.DataFrame:
         A pandas dataframe showing the accessible dataflow storage accounts.
     """
 
-    df = pd.DataFrame(
-        columns=[
-            "Dataflow Storage Account ID",
-            "Dataflow Storage Account Name",
-            "Enabled",
-        ]
-    )
-    client = fabric.PowerBIRestClient()
-    response = client.get("/v1.0/myorg/dataflowStorageAccounts")
-    if response.status_code != 200:
-        raise FabricHTTPException(response)
+    columns = {
+        "Dataflow Storage Account ID": "string",
+        "Dataflow Storage Account Name": "string",
+        "Enabled": "bool",
+    }
+    df = _create_dataframe(columns=columns)
+
+    response = _base_api(request="/v1.0/myorg/dataflowStorageAccounts")
 
     for v in response.json().get("value", []):
         new_data = {
@@ -132,7 +133,7 @@ def list_dataflow_storage_accounts() -> pd.DataFrame:
         }
         df = pd.concat([df, pd.DataFrame(new_data, index=[0])], ignore_index=True)
 
-    df["Enabled"] = df["Enabled"].astype(bool)
+    _update_dataframe_datatypes(dataframe=df, column_map=columns)
 
     return df
 
@@ -164,29 +165,23 @@ def list_upstream_dataflows(
     (dataflow_name, dataflow_id) = _resolve_dataflow_name_and_id(
         dataflow=dataflow, workspace=workspace_id
     )
-    client = fabric.PowerBIRestClient()
-
-    df = pd.DataFrame(
-        columns=[
-            "Dataflow Name",
-            "Dataflow Id",
-            "Workspace Name",
-            "Workspace Id",
-            "Upstream Dataflow Name",
-            "Upstream Dataflow Id",
-            "Upstream Workspace Name",
-            "Upstream Workspace Id",
-        ]
-    )
 
-    def collect_upstreams(
-        client, dataflow_id, dataflow_name, workspace_id, workspace_name
-    ):
-        response = client.get(
-            f"/v1.0/myorg/groups/{workspace_id}/dataflows/{dataflow_id}/upstreamDataflows"
+    columns = {
+        "Dataflow Name": "string",
+        "Dataflow Id": "string",
+        "Workspace Name": "string",
+        "Workspace Id": "string",
+        "Upstream Dataflow Name": "string",
+        "Upstream Dataflow Id": "string",
+        "Upstream Workspace Name": "string",
+        "Upstream Workspace Id": "string",
+    }
+    df = _create_dataframe(columns=columns)
+
+    def collect_upstreams(dataflow_id, dataflow_name, workspace_id, workspace_name):
+        response = _base_api(
+            request=f"/v1.0/myorg/groups/{workspace_id}/dataflows/{dataflow_id}/upstreamDataflows"
         )
-        if response.status_code != 200:
-            raise FabricHTTPException(response)
 
         values = response.json().get("value", [])
         for v in values:
@@ -209,14 +204,13 @@ def list_upstream_dataflows(
             }
 
             collect_upstreams(
-                client,
                 tgt_dataflow_id,
                 tgt_dataflow_name,
                 tgt_workspace_id,
                 tgt_workspace_name,
             )
 
-    collect_upstreams(client, dataflow_id, dataflow_name, workspace_id, workspace_name)
+    collect_upstreams(dataflow_id, dataflow_name, workspace_id, workspace_name)
 
     return df
 

sempy_labs/_dax.py

@@ -4,12 +4,16 @@ from sempy_labs._helper_functions import (
     resolve_workspace_name_and_id,
     format_dax_object_name,
     resolve_dataset_name_and_id,
+    _base_api,
+    generate_guid,
 )
 from sempy_labs._model_dependencies import get_model_calc_dependencies
-from typing import Optional, List
+from typing import Optional, List, Tuple
 from sempy._utils._log import log
 from uuid import UUID
 from sempy_labs.directlake._warm_cache import _put_columns_into_memory
+import sempy_labs._icons as icons
+import time
 
 
 @log
@@ -47,15 +51,15 @@ def evaluate_dax_impersonation(
     (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
     (dataset_name, dataset_id) = resolve_dataset_name_and_id(dataset, workspace_id)
 
-    request_body = {
+    payload = {
         "queries": [{"query": dax_query}],
         "impersonatedUserName": user_name,
     }
 
-    client = fabric.PowerBIRestClient()
-    response = client.post(
-        f"/v1.0/myorg/groups/{workspace_id}/datasets/{dataset_id}/executeQueries",
-        json=request_body,
+    response = _base_api(
+        request=f"/v1.0/myorg/groups/{workspace_id}/datasets/{dataset_id}/executeQueries",
+        method="post",
+        payload=payload,
     )
     data = response.json()["results"][0]["tables"]
     column_names = data[0]["rows"][0].keys()
@@ -100,6 +104,8 @@ def get_dax_query_dependencies(
     (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
     (dataset_name, dataset_id) = resolve_dataset_name_and_id(dataset, workspace_id)
 
+    fabric.refresh_tom_cache(workspace=workspace)
+
     if isinstance(dax_string, str):
         dax_string = [dax_string]
 
@@ -257,3 +263,219 @@ def get_dax_query_memory_size(
     )
 
     return df["Total Size"].sum()
+
+
+@log
+def _dax_perf_test(
+    dataset: str,
+    dax_queries: dict,
+    clear_cache_before_run: bool = False,
+    refresh_type: Optional[str] = None,
+    rest_time: int = 2,
+    workspace: Optional[str] = None,
+) -> Tuple[pd.DataFrame, dict]:
+    """
+    Runs a performance test on a set of DAX queries.
+
+    Parameters
+    ----------
+    dataset : str
+        Name of the semantic model.
+    dax_queries : dict
+        The dax queries to run in a dictionary format. Here is an example:
+        {
+            "Sales Amount Test", """ """ EVALUATE SUMMARIZECOLUMNS("Sales Amount", [Sales Amount]) """ """,
+            "Order Quantity with Product", """ """ EVALUATE SUMMARIZECOLUMNS('Product'[Color], "Order Qty", [Order Qty]) """ """,
+        }
+    clear_cache_before_run : bool, default=False
+    refresh_type : str, default=None
+    rest_time : int, default=2
+        Rest time (in seconds) between the execution of each DAX query.
+    workspace : str, default=None
+        The Fabric workspace name.
+        Defaults to None which resolves to the workspace of the attached lakehouse
+        or if no lakehouse attached, resolves to the workspace of the notebook.
+
+    Returns
+    -------
+    Tuple[pandas.DataFrame, dict]
+        A pandas dataframe showing the SQL profiler trace results of the DAX queries.
+        A dictionary of the query results in pandas dataframes.
+    """
+    from sempy_labs._refresh_semantic_model import refresh_semantic_model
+    from sempy_labs._clear_cache import clear_cache
+
+    event_schema = {
+        "QueryBegin": [
+            "EventClass",
+            "EventSubclass",
+            "CurrentTime",
+            "NTUserName",
+            "TextData",
+            "StartTime",
+            "ApplicationName",
+        ],
+        "QueryEnd": [
+            "EventClass",
+            "EventSubclass",
+            "CurrentTime",
+            "NTUserName",
+            "TextData",
+            "StartTime",
+            "EndTime",
+            "Duration",
+            "CpuTime",
+            "Success",
+            "ApplicationName",
+        ],
+        "VertiPaqSEQueryBegin": [
+            "EventClass",
+            "EventSubclass",
+            "CurrentTime",
+            "NTUserName",
+            "TextData",
+            "StartTime",
+        ],
+        "VertiPaqSEQueryEnd": [
+            "EventClass",
+            "EventSubclass",
+            "CurrentTime",
+            "NTUserName",
+            "TextData",
+            "StartTime",
+            "EndTime",
+            "Duration",
+            "CpuTime",
+            "Success",
+        ],
+        "VertiPaqSEQueryCacheMatch": [
+            "EventClass",
+            "EventSubclass",
+            "CurrentTime",
+            "NTUserName",
+            "TextData",
+        ],
+    }
+
+    # Add Execution Metrics
+    event_schema["ExecutionMetrics"] = ["EventClass", "ApplicationName", "TextData"]
+    # Add DAX Query Plan
+    # event_schema["DAXQueryPlan"] = ["EventClass", "EventSubclass", "CurrentTime", "StartTime", "EndTime", "Duration", "CpuTime", "ApplicationName", "TextData"]
+
+    query_results = {}
+
+    # Establish trace connection
+    with fabric.create_trace_connection(
+        dataset=dataset, workspace=workspace
+    ) as trace_connection:
+        with trace_connection.create_trace(event_schema) as trace:
+            trace.start()
+            print(f"{icons.in_progress} Starting performance testing...")
+            # Loop through DAX queries
+            for name, dax in dax_queries.items():
+
+                if clear_cache_before_run:
+                    clear_cache(dataset=dataset, workspace=workspace)
+                if refresh_type is not None:
+                    refresh_semantic_model(
+                        dataset=dataset, workspace=workspace, refresh_type=refresh_type
+                    )
+
+                # EVALUATE {1} is used to initate a warm cache
+                fabric.evaluate_dax(
+                    dataset=dataset, workspace=workspace, dax_string="""EVALUATE {1}"""
+                )
+                # Run DAX Query
+                result = fabric.evaluate_dax(
+                    dataset=dataset, workspace=workspace, dax_string=dax
+                )
+
+                # Add results to output
+                query_results[name] = result
+
+                time.sleep(rest_time)
+                print(f"{icons.green_dot} The '{name}' query has completed.")
+
+            df = trace.stop()
+            # Allow time to collect trace results
+            time.sleep(5)
+
+            # Step 1: Filter out unnecessary operations
+            query_names = list(dax_queries.keys())
+            df = df[
+                ~df["Application Name"].isin(["PowerBI", "PowerBIEIM"])
+                & (~df["Text Data"].str.startswith("EVALUATE {1}"))
+            ]
+            query_begin = df["Event Class"] == "QueryBegin"
+            temp_column_name = "QueryName_INT"
+            df = df.copy()
+            df[temp_column_name] = query_begin.cumsum()
+            df[temp_column_name] = (
+                df[temp_column_name]
+                .where(query_begin, None)  # Assign None to non-query begin rows
+                .ffill()  # Forward fill None values
+                .astype("Int64")  # Use pandas nullable integer type for numeric indices
+            )
+
+            df.loc[df[temp_column_name].notna(), "Query Name"] = (
+                df[temp_column_name]
+                .dropna()
+                .astype(int)
+                .map(lambda x: query_names[x - 1])
+            )
+            df = df[df[temp_column_name] != None]
+            df = df.drop(columns=[temp_column_name])
+
+            query_to_guid = {
+                name: generate_guid() for name in df["Query Name"].unique()
+            }
+            df["Query ID"] = df["Query Name"].map(query_to_guid)
+
+    df = df.reset_index(drop=True)
+
+    return df, query_results
+
+
+def _dax_perf_test_bulk(
+    mapping: dict,
+    clear_cache_before_run: bool = False,
+    refresh_type: Optional[str] = None,
+    rest_time: int = 2,
+):
+    """
+    mapping is something like this:
+
+    mapping = {
+        "Workspace1": {
+            "Dataset1": {
+                "Query1": "EVALUATE ...",
+                "Query2": "EVALUATE ...",
+            },
+            "Dataset2": {
+                "Query3": "EVALUATE ...",
+                "Query4": "EVALUATE ...",
+            }
+        },
+        "Workspace2": {
+            "Dataset3": {
+                "Query5": "EVALUATE ...",
+                "Query6": "EVALUATE ...",
+            },
+            "Dataset4": {
+                "Query7": "EVALUATE ...",
+                "Query8": "EVALUATE ...",
+            }
+        }
+    }
+    """
+
+    for workspace, datasets in mapping.items():
+        for dataset, queries in datasets.items():
+            _dax_perf_test(
+                dataset=dataset,
+                dax_queries=queries,
+                clear_cache_before_run=clear_cache_before_run,
+                refresh_type=refresh_type,
+                rest_time=rest_time,
+                workspace=workspace,
+            )