moved logic to lro function

Author: m-kovalsky

src/sempy_labs/_helper_functions.py

@@ -1568,7 +1568,9 @@ def lro(
     status_codes: Optional[List[str]] = [200, 202],
     sleep_time: Optional[int] = 1,
     return_status_code: bool = False,
+    job_scheduler: bool = False,
 ):
+    from sempy_labs._job_scheduler import _get_item_job_instance
 
     if response.status_code not in status_codes:
         raise FabricHTTPException(response)
@@ -1578,20 +1580,32 @@ def lro(
         else:
             result = response
     if response.status_code == status_codes[1]:
-        operationId = response.headers["x-ms-operation-id"]
-        response = client.get(f"/v1/operations/{operationId}")
-        response_body = json.loads(response.content)
-        while response_body["status"] not in ["Succeeded", "Failed"]:
-            time.sleep(sleep_time)
-            response = client.get(f"/v1/operations/{operationId}")
-            response_body = json.loads(response.content)
-        if response_body["status"] != "Succeeded":
-            raise FabricHTTPException(response)
-        if return_status_code:
-            result = response.status_code
+        if job_scheduler:
+            status_url = response.headers.get("Location").split("fabric.microsoft.com")[
+                1
+            ]
+            status = None
+            while status not in ["Completed", "Failed"]:
+                response = _base_api(request=status_url)
+                status = response.json().get("status")
+                time.sleep(3)
+
+            return _get_item_job_instance(url=status_url)
         else:
-            response = client.get(f"/v1/operations/{operationId}/result")
-            result = response
+            operation_id = response.headers["x-ms-operation-id"]
+            response = client.get(f"/v1/operations/{operation_id}")
+            response_body = json.loads(response.content)
+            while response_body["status"] not in ["Succeeded", "Failed"]:
+                time.sleep(sleep_time)
+                response = client.get(f"/v1/operations/{operation_id}")
+                response_body = json.loads(response.content)
+            if response_body["status"] != "Succeeded":
+                raise FabricHTTPException(response)
+            if return_status_code:
+                result = response.status_code
+            else:
+                response = client.get(f"/v1/operations/{operation_id}/result")
+                result = response
 
     return result
 
@@ -2212,6 +2226,7 @@ def _base_api(
     uses_pagination: bool = False,
     lro_return_json: bool = False,
     lro_return_status_code: bool = False,
+    lro_return_df: bool = False,
 ):
     import notebookutils
     from sempy_labs._authentication import _get_headers
@@ -2267,7 +2282,9 @@ def get_token(self, *scopes, **kwargs) -> AccessToken:
             json=payload,
         )
 
-    if lro_return_json:
+    if lro_return_df:
+        return lro(c, response, status_codes, job_scheduler=True)
+    elif lro_return_json:
         return lro(c, response, status_codes).json()
     elif lro_return_status_code:
         return lro(c, response, status_codes, return_status_code=True)

src/sempy_labs/graph_model/_background_jobs.py

@@ -8,8 +8,6 @@
     resolve_workspace_name_and_id,
 )
 import sempy_labs._icons as icons
-import time
-from sempy_labs._job_scheduler import _get_item_job_instance
 
 
 @log
@@ -41,23 +39,16 @@ def refresh_graph(
         item=graph_model, type="GraphModel", workspace=workspace_id
     )
 
-    response = _base_api(
-        request=f"/v1/workspaces/{workspace_id}/GraphModels/{item_id}/jobs/instances?jobType=RefreshGraph",
-        method="post",
-    )
-
     print(
         f"{icons.in_progress} The refresh graph job for the '{item_name}' graph model within the '{workspace_name}' workspace has been initiated."
     )
 
-    status_url = response.headers.get("Location").split("fabric.microsoft.com")[1]
-    status = None
-    while status not in ["Completed", "Failed"]:
-        response = _base_api(request=status_url)
-        status = response.json().get("status")
-        time.sleep(3)
-
-    df = _get_item_job_instance(url=status_url)
+    df = _base_api(
+        request=f"/v1/workspaces/{workspace_id}/GraphModels/{item_id}/jobs/instances?jobType=RefreshGraph",
+        method="post",
+        lro_return_df=True,
+    )
+    status = df["Status"].iloc[0]
 
     if status == "Completed":
         print(

src/sempy_labs/lakehouse/_lakehouse.py

@@ -11,11 +11,7 @@
 )
 import sempy_labs._icons as icons
 import re
-import time
 import pandas as pd
-from sempy_labs._job_scheduler import (
-    _get_item_job_instance,
-)
 
 
 @log
@@ -247,26 +243,20 @@ def is_valid_format(time_string):
     if vacuum and retention_period is not None:
         payload["executionData"]["vacuumSettings"]["retentionPeriod"] = retention_period
 
-    response = _base_api(
+    print(
+        f"{icons.in_progress} The table maintenance job for the '{table_name}' table in the '{lakehouse_name}' lakehouse within the '{workspace_name}' workspace has been initiated."
+    )
+
+    df = _base_api(
         request=f"/v1/workspaces/{workspace_id}/lakehouses/{lakehouse_id}/jobs/instances?jobType=TableMaintenance",
         method="post",
         payload=payload,
-        status_codes=202,
+        status_codes=[200, 202],
         client="fabric_sp",
+        lro_return_df=True,
     )
 
-    print(
-        f"{icons.in_progress} The table maintenance job for the '{table_name}' table in the '{lakehouse_name}' lakehouse within the '{workspace_name}' workspace has been initiated."
-    )
-
-    status_url = response.headers.get("Location").split("fabric.microsoft.com")[1]
-    status = None
-    while status not in ["Completed", "Failed"]:
-        response = _base_api(request=status_url)
-        status = response.json().get("status")
-        time.sleep(3)
-
-    df = _get_item_job_instance(url=status_url)
+    status = df["Status"].iloc[0]
 
     if status == "Completed":
         print(

src/sempy_labs/lakehouse/_materialized_lake_views.py

@@ -9,10 +9,6 @@
 from uuid import UUID
 from sempy._utils._log import log
 import sempy_labs._icons as icons
-import time
-from sempy_labs._job_scheduler import (
-    _get_item_job_instance,
-)
 import pandas as pd
 
 
@@ -46,22 +42,16 @@ def refresh_materialized_lake_views(
         lakehouse=lakehouse, workspace=workspace_id
     )
 
-    response = _base_api(
-        request=f"/v1/workspaces/{workspace_id}/lakehouses/{lakehouse_id}/jobs/instances?jobType=RefreshMaterializedLakeViews",
-    )
-
     print(
         f"{icons.in_progress} The refresh materialized lake views job for the '{lakehouse_name}' lakehouse within the '{workspace_name}' workspace has been initiated."
     )
 
-    status_url = response.headers.get("Location").split("fabric.microsoft.com")[1]
-    status = None
-    while status not in ["Completed", "Failed"]:
-        response = _base_api(request=status_url)
-        status = response.json().get("status")
-        time.sleep(3)
+    df = _base_api(
+        request=f"/v1/workspaces/{workspace_id}/lakehouses/{lakehouse_id}/jobs/instances?jobType=RefreshMaterializedLakeViews",
+        lro_return_df=True,
+    )
 
-    df = _get_item_job_instance(url=status_url)
+    status = df["Status"].iloc[0]
 
     if status == "Completed":
         print(