black + gen dl

Author: m-kovalsky

src/sempy_labs/directlake/_dl_helper.py

@@ -6,6 +6,7 @@
 import sempy_labs._icons as icons
 from sempy._utils._log import log
 from sempy_labs._helper_functions import (
+    resolve_lakehouse_name_and_id,
     retry,
     _convert_data_type,
     resolve_dataset_name_and_id,
@@ -83,7 +84,7 @@ def generate_direct_lake_semantic_model(
     workspace: Optional[str | UUID] = None,
     lakehouse: Optional[str] = None,
     lakehouse_workspace: Optional[str | UUID] = None,
-    schema: str = "dbo",
+    schema: Optional[str] = None,
     overwrite: bool = False,
     refresh: bool = True,
 ):
@@ -107,7 +108,7 @@ def generate_direct_lake_semantic_model(
         The Fabric workspace name or ID in which the lakehouse resides.
         Defaults to None which resolves to the workspace of the attached lakehouse
         or if no lakehouse attached, resolves to the workspace of the notebook.
-    schema : str, default="dbo"
+    schema : str, default=None
         The schema used for the lakehouse.
     overwrite : bool, default=False
         If set to True, overwrites the existing semantic model if it already exists.
@@ -127,10 +128,16 @@ def generate_direct_lake_semantic_model(
         lakehouse_tables = [lakehouse_tables]
 
     (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
-    if lakehouse_workspace is None:
-        lakehouse_workspace = workspace
+    (lakehouse_workspace_name, lakehouse_workspace_id) = resolve_workspace_name_and_id(
+        lakehouse_workspace
+    )
+    (lakehouse_name, lakehouse_id) = resolve_lakehouse_name_and_id(
+        lakehouse, lakehouse_workspace_id
+    )
 
-    dfLT = get_lakehouse_tables(lakehouse=lakehouse, workspace=lakehouse_workspace)
+    dfLT = get_lakehouse_tables(
+        lakehouse=lakehouse_id, workspace=lakehouse_workspace_id
+    )
 
     icons.sll_tags.append("GenerateDLModel")
 
@@ -141,9 +148,9 @@ def generate_direct_lake_semantic_model(
                 f"{icons.red_dot} The '{t}' table does not exist as a delta table in the '{lakehouse}' within the '{workspace_name}' workspace."
             )
 
-    dfLC = get_lakehouse_columns(lakehouse=lakehouse, workspace=lakehouse_workspace)
+    dfLC = get_lakehouse_columns(lakehouse=lakehouse, workspace=lakehouse_workspace_id)
     expr = generate_shared_expression(
-        item_name=lakehouse, item_type="Lakehouse", workspace=lakehouse_workspace
+        item_name=lakehouse, item_type="Lakehouse", workspace=lakehouse_workspace_id
     )
     dfD = fabric.list_datasets(workspace=workspace_id, mode="rest")
     dfD_filt = dfD[dfD["Dataset Name"] == dataset]
@@ -179,9 +186,17 @@ def dyn_connect():
 
         for t in lakehouse_tables:
             tom.add_table(name=t)
-            tom.add_entity_partition(table_name=t, entity_name=t, schema_name=schema)
+            tom.add_entity_partition(
+                table_name=t, entity_name=t, schema_name=schema or "dbo"
+            )
             dfLC_filt = dfLC[dfLC["Table Name"] == t]
-            for i, r in dfLC_filt.iterrows():
+            if schema:
+                dfLC_filt = dfLC_filt[dfLC_filt["Schema Name"] == schema]
+            if dfLC_filt.empty:
+                raise ValueError(
+                    f"{icons.red_dot} No columns found for the '{t}' table in the '{lakehouse_name}' lakehouse. Please ensure the table has columns and if a schema is provided, that the schema name is correct."
+                )
+            for _, r in dfLC_filt.iterrows():
                 lakeCName = r["Column Name"]
                 dType = r["Data Type"]
                 dt = _convert_data_type(dType)

src/sempy_labs/directlake/_update_directlake_partition_entity.py

@@ -86,9 +86,7 @@ def update_direct_lake_partition_entity(
             # Update source lineage tag
             if schema:
                 # Updated to all DL patterns (DLOL, DLSQL)
-                tom.model.Tables[tName].Partitions[
-                        part_name
-                    ].Source.SchemaName = schema
+                tom.model.Tables[tName].Partitions[part_name].Source.SchemaName = schema
                 tom.model.Tables[tName].SourceLineageTag = f"[{schema}].[{eName}]"
             else:
                 tom.model.Tables[tName].SourceLineageTag = f"[dbo].[{eName}]"

src/sempy_labs/lakehouse/_get_lakehouse_columns.py

@@ -78,7 +78,7 @@ def add_column_metadata(table_name, schema_name, col_name, data_type):
 
     for _, r in tables_filt.iterrows():
         table_name = r["Table Name"]
-        schema_name = r['Schema Name']
+        schema_name = r["Schema Name"]
         path = r["Location"]
 
         if _pure_python_notebook():

src/sempy_labs/report/_reportwrapper.py

@@ -100,8 +100,11 @@ def __init__(
         self._show_diffs = show_diffs
 
         # Check that the report is in the PBIR format.
-        response = _base_api(request=f"/v1.0/myorg/groups/{self._workspace_id}/reports/{self._report_id}", client="fabric_sp")
-        self.format = response.json().get('format')
+        response = _base_api(
+            request=f"/v1.0/myorg/groups/{self._workspace_id}/reports/{self._report_id}",
+            client="fabric_sp",
+        )
+        self.format = response.json().get("format")
 
         result = _base_api(
             request=f"/v1/workspaces/{self._workspace_id}/items/{self._report_id}/getDefinition",

src/sempy_labs/semantic_model/_caching.py

@@ -12,7 +12,9 @@
 
 @log
 def enable_query_caching(
-    dataset: str | UUID, workspace: Optional[str | UUID] = None, enable: bool = True,
+    dataset: str | UUID,
+    workspace: Optional[str | UUID] = None,
+    enable: bool = True,
 ):
     """
     Enables or disables `query caching <http://aka.ms/queryCaching>`_ for a semantic model.
@@ -35,7 +37,9 @@ def enable_query_caching(
     )
     model_id = get_model_id(item_id=item_id)
     if model_id is None:
-        raise ValueError(f"Failed to retrieve model ID for semantic model '{item_name}'")
+        raise ValueError(
+            f"Failed to retrieve model ID for semantic model '{item_name}'"
+        )
 
     caching_map = {
         True: 2,
@@ -44,7 +48,12 @@ def enable_query_caching(
 
     payload = {"queryCachingState": caching_map.get(enable)}
 
-    _base_api(request=f"metadata/models/{model_id}/caching", method="internal", payload=payload, status_codes=204)
+    _base_api(
+        request=f"metadata/models/{model_id}/caching",
+        method="internal",
+        payload=payload,
+        status_codes=204,
+    )
 
     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/surge_protection/_items.py

@@ -10,13 +10,17 @@
 from sempy._utils._log import log
 
 
-def _surge_api(
-    capacity, url, payload, method="get", status_code=200, return_json=True
-):
+def _surge_api(capacity, url, payload, method="get", status_code=200, return_json=True):
 
     capacity_id = resolve_capacity_id(capacity)
 
-    response = _base_api(request=f"capacities/{capacity_id}/{url}", client="internal", method=method, payload=payload, status_codes=status_code)
+    response = _base_api(
+        request=f"capacities/{capacity_id}/{url}",
+        client="internal",
+        method=method,
+        payload=payload,
+        status_codes=status_code,
+    )
 
     if return_json:
         return response.json()
@@ -48,9 +52,7 @@ def get_workspace_consumption_rules(
         or a dictionary if return_dataframe is set to False.
     """
 
-    response_json = _surge_api(
-        capacity=capacity, url="detectionRules", payload=None
-    )
+    response_json = _surge_api(capacity=capacity, url="detectionRules", payload=None)
 
     if not return_dataframe:
         return response_json
@@ -347,4 +349,6 @@ def delete_background_operation_rules(capacity: str | UUID = None):
         return_json=False,
     )
 
-    print(f"{icons.green_dot} The background operation rules have been deleted successfully.")
+    print(
+        f"{icons.green_dot} The background operation rules have been deleted successfully."
+    )