Notebook with sempy_labs failed when executed in DataPipeline

I have implemented a workaround function to bypass editing the XMLA endpoint (directlake.update_direct_lake_model_connection()) using a REST-based approach.

Here is my function to update the direct lake connection

def update_direct_lake_onelake_connection(
    dataset_name: str,
    workspace_name: str,
    lakehouse_name: str,
    lakehouse_workspace_name: str,
) -> None:
    """
    Patch the OneLake DFS URL in a Direct Lake semantic model's shared
    expressions to point to the correct target Lakehouse.

    Uses the Fabric REST API (getDefinition / updateDefinition) so that
    Service Principal authentication is honoured correctly — bypassing
    the XMLA endpoint where SP tokens are not forwarded by sempy_labs.

    Parameters
    ----------
    dataset_name             : Display name of the target semantic model.
    workspace_name           : Display name of the workspace containing the model.
    lakehouse_name           : Display name of the target lakehouse.
    lakehouse_workspace_name : Display name of the workspace containing the lakehouse.
    """
    logger.info("Resolving IDs for Direct Lake connection update...")

    workspace_id          = get_workspace_id(workspace_name)
    lakehouse_workspace_id = get_workspace_id(lakehouse_workspace_name)
    lakehouse_id          = get_lakehouse_id(lakehouse_workspace_id, lakehouse_name)
    dataset_id            = get_semantic_model_id(workspace_id, dataset_name)

    if not dataset_id:
        raise ValueError(f"Semantic model '{dataset_name}' not found in workspace '{workspace_name}'.")

    logger.info(f"  Workspace ID:           {workspace_id}")
    logger.info(f"  Lakehouse Workspace ID: {lakehouse_workspace_id}")
    logger.info(f"  Lakehouse ID:           {lakehouse_id}")
    logger.info(f"  Semantic Model ID:      {dataset_id}")

    # --------------------------------------------------
    # Get model definition (handles sync 200 and async 202)
    # --------------------------------------------------
    r = requests.post(
        f"{FABRIC_API_BASE}/workspaces/{workspace_id}/semanticModels/{dataset_id}/getDefinition",
        headers=HEADERS,
    )
    r.raise_for_status()

    if r.status_code == 202:
        poll_result = poll_operation(r.headers["Location"], "getDefinition")
        if "definition" in poll_result:
            parts = poll_result["definition"]["parts"]
        else:
            result_url = poll_result.get("resourceLocation") or r.headers["Location"] + "/result"
            result_r = requests.get(result_url, headers=HEADERS)
            result_r.raise_for_status()
            parts = result_r.json()["definition"]["parts"]
    else:
        parts = r.json()["definition"]["parts"]

    # --------------------------------------------------
    # Patch OneLake URL in all matching parts
    # --------------------------------------------------
    new_url = (
        f"https://onelake.dfs.fabric.microsoft.com/"
        f"{lakehouse_workspace_id}/{lakehouse_id}"
    )
    onelake_pattern = re.compile(
        r"https://onelake\.dfs\.fabric\.microsoft\.com/"
        r"[0-9a-fA-F\-]{36}/[0-9a-fA-F\-]{36}"
    )
    patched = False

    for part in parts:
        content = base64.b64decode(part["payload"]).decode("utf-8")

        if "onelake.dfs.fabric.microsoft.com" not in content.lower():
            continue

        updated = onelake_pattern.sub(new_url, content)

        if updated == content:
            logger.warning(
                f"Part '{part['path']}' contains OneLake URL but regex did not match — "
                f"inspect manually:\n{content}"
            )
            continue

        logger.info(f"  Patched part: {part['path']}")
        part["payload"] = base64.b64encode(updated.encode("utf-8")).decode("utf-8")
        patched = True

    if not patched:
        raise ValueError(
            f"No patchable OneLake URL found in model '{dataset_name}'. "
            "Verify this is a Direct Lake model."
        )

    # --------------------------------------------------
    # Push updated definition back (handles sync / async)
    # --------------------------------------------------
    r2 = requests.post(
        f"{FABRIC_API_BASE}/workspaces/{workspace_id}/semanticModels/{dataset_id}/updateDefinition",
        headers=HEADERS,
        json={"definition": {"parts": parts}},
    )
    r2.raise_for_status()

    if r2.status_code == 202:
        poll_operation(r2.headers["Location"], "updateDefinition")

    logger.info("Direct Lake OneLake connection updated successfully.")