feat: Add save_failed_payload debugging flag to call-endpoint functions

Author: Thijs Lukkezen

src/pypilecore/api.py

@@ -1,6 +1,6 @@
 import logging
 from time import sleep
-from typing import Literal
+from typing import Literal, overload
 
 from nuclei.client import NucleiClient
 from requests import Response
@@ -69,8 +69,90 @@ def wait_until_ticket_is_ready(
             raise RuntimeError(failure_message)
 
 
+@overload
+def get_task_result_pipeline(
+    client: NucleiClient,
+    endpoint: str,
+    payload: dict,
+    verbose: bool,
+    response_type: Literal["dict"],
+    save_failed_payload: bool,
+    failed_payload_filename: str,
+) -> dict: ...
+
+
+@overload
+def get_task_result_pipeline(
+    client: NucleiClient,
+    endpoint: str,
+    payload: dict,
+    verbose: bool,
+    response_type: Literal["bytes"],
+    save_failed_payload: bool,
+    failed_payload_filename: str,
+) -> bytes: ...
+
+
+def get_task_result_pipeline(
+    client: NucleiClient,
+    endpoint: str,
+    payload: dict,
+    verbose: bool = False,
+    response_type: Literal["dict", "bytes"] = "dict",
+    save_failed_payload: bool = False,
+    failed_payload_filename: str = "debug_payload.json",
+) -> dict | bytes:
+    """
+    Execute a task and return the result.
+
+    Parameters
+    ----------
+    client: NucleiClient
+        client object created by [nuclei](https://github.com/cemsbv/nuclei)
+    endpoint: str
+        the endpoint to call
+    payload: dict
+        the payload of the request
+    verbose: bool
+        if True, print additional information to the console
+    save_failed_payload: bool
+        if True, save the payload to a file in case of failure
+    failed_payload_filename: str
+        the name of the file to save the failed payload to
+    """
+
+    try:
+        ticket = client.call_endpoint(
+            "PileCore",
+            endpoint,
+            version="v4",
+            schema=payload,
+            return_response=True,
+        )
+
+        wait_until_ticket_is_ready(client=client, ticket=ticket, verbose=verbose)
+
+        response = client.call_endpoint(
+            "PileCore", "/get-task-result", version="v4", schema=ticket.json()
+        )
+    except Exception as e:
+        if save_failed_payload:
+            # In case of failure, save the payload to a file for debugging
+            with open(failed_payload_filename, "w") as f:
+                from nuclei.client.utils import serialize_json_string
+
+                f.write(serialize_json_string(payload))
+        raise e
+
+    return response
+
+
 def get_multi_cpt_api_result(
-    client: NucleiClient, payload: dict, verbose: bool = False
+    client: NucleiClient,
+    payload: dict,
+    verbose: bool = False,
+    save_failed_payload: bool = False,
+    failed_payload_filename: str = "pilecore_multi_cpt_result_debug_payload.json",
 ) -> dict:
     """
     Wrapper around the PileCore endpoint "/bearing/multiple-cpts/results".
@@ -88,24 +170,24 @@ def get_multi_cpt_api_result(
         "Calculating bearing capacities... \n"
         "Depending on the amount of pile tip levels and CPT's this can take a while."
     )
-    ticket = client.call_endpoint(
-        "PileCore",
-        "/bearing/multiple-cpts/results",
-        version="v4",
-        schema=payload,
-        return_response=True,
-    )
-
-    wait_until_ticket_is_ready(client=client, ticket=ticket, verbose=verbose)
-
-    return client.call_endpoint(
-        "PileCore", "/get-task-result", version="v4", schema=ticket.json()
+    return get_task_result_pipeline(
+        client=client,
+        endpoint="/bearing/multiple-cpts/results",
+        payload=payload,
+        verbose=verbose,
+        response_type="dict",
+        save_failed_payload=save_failed_payload,
+        failed_payload_filename=failed_payload_filename,
     )
 
 
 def get_multi_cpt_api_report(
-    client: NucleiClient, payload: dict, verbose: bool = False
-) -> dict:
+    client: NucleiClient,
+    payload: dict,
+    verbose: bool = False,
+    save_failed_payload: bool = False,
+    failed_payload_filename: str = "pilecore_multi_cpt_report_debug_payload.json",
+) -> bytes:
     """
     Wrapper around the PileCore endpoint "/bearing/multiple-cpts/report".
 
@@ -122,22 +204,23 @@ def get_multi_cpt_api_report(
         "Generate report... \n"
         "Depending on the amount of pile tip levels and CPT's this can take a while."
     )
-    ticket = client.call_endpoint(
-        "PileCore",
-        "/bearing/multiple-cpts/report",
-        version="v4",
-        schema=payload,
-        return_response=True,
-    )
-    wait_until_ticket_is_ready(client=client, ticket=ticket, verbose=verbose)
-
-    return client.call_endpoint(
-        "PileCore", "/get-task-result", version="v4", schema=ticket.json()
+    return get_task_result_pipeline(
+        client=client,
+        endpoint="/bearing/multiple-cpts/report",
+        payload=payload,
+        verbose=verbose,
+        response_type="bytes",
+        save_failed_payload=save_failed_payload,
+        failed_payload_filename=failed_payload_filename,
     )
 
 
 def get_groups_api_result(
-    client: NucleiClient, payload: dict, verbose: bool = False
+    client: NucleiClient,
+    payload: dict,
+    verbose: bool = False,
+    save_failed_payload: bool = False,
+    failed_payload_filename: str = "pilecore_group_cpts_debug_payload.json",
 ) -> dict:
     """
     Wrapper around the PileCore endpoint "/grouper/group_cpts".
@@ -155,23 +238,23 @@ def get_groups_api_result(
         "Finding groups... \n"
         "Depending on the amount of pile tip levels and CPT's this can take a while."
     )
-    ticket = client.call_endpoint(
-        "PileCore",
-        "/grouper/group-cpts",
-        version="v4",
-        schema=payload,
-        return_response=True,
-    )
-
-    wait_until_ticket_is_ready(client=client, ticket=ticket, verbose=verbose)
-
-    return client.call_endpoint(
-        "PileCore", "/get-task-result", version="v4", schema=ticket.json()
+    return get_task_result_pipeline(
+        client=client,
+        endpoint="/grouper/group-cpts",
+        payload=payload,
+        verbose=verbose,
+        response_type="dict",
+        save_failed_payload=save_failed_payload,
+        failed_payload_filename=failed_payload_filename,
     )
 
 
 def get_groups_api_report(
-    client: NucleiClient, payload: dict, verbose: bool = False
+    client: NucleiClient,
+    payload: dict,
+    verbose: bool = False,
+    save_failed_payload: bool = False,
+    failed_payload_filename: str = "pilecore_grouper_report_debug_payload.json",
 ) -> bytes:
     """
     Wrapper around the PileCore endpoint "/grouper/generate_grouper_report".
@@ -189,18 +272,14 @@ def get_groups_api_report(
         "Generate report... \n"
         "Depending on the amount of pile tip levels and CPT's this can take a while."
     )
-    ticket = client.call_endpoint(
-        "PileCore",
-        "/grouper/generate-grouper-report",
-        version="v4",
-        schema=payload,
-        return_response=True,
-    )
-
-    wait_until_ticket_is_ready(client=client, ticket=ticket, verbose=verbose)
-
-    return client.call_endpoint(
-        "PileCore", "/get-task-result", version="v4", schema=ticket.json()
+    return get_task_result_pipeline(
+        client=client,
+        endpoint="/grouper/generate-grouper-report",
+        payload=payload,
+        verbose=verbose,
+        response_type="bytes",
+        save_failed_payload=save_failed_payload,
+        failed_payload_filename=failed_payload_filename,
     )
 
 
@@ -209,6 +288,8 @@ def get_multi_cpt_api_result_tension(
     payload: dict,
     verbose: bool = False,
     standard: Literal["NEN9997-1", "CUR236"] = "NEN9997-1",
+    save_failed_payload: bool = False,
+    failed_payload_filename: str = "pilecore_multi_cpt_tension_result_debug_payload.json",
 ) -> dict:
     """
     Wrapper around the PileCore endpoint "/uplift/[nen or cur]/multiple-cpts/results".
@@ -234,18 +315,14 @@ def get_multi_cpt_api_result_tension(
         endpoint = "/uplift/cur/multiple-cpts/results"
         payload.pop("construction_sequence", None)
 
-    ticket = client.call_endpoint(
-        "PileCore",
+    return get_task_result_pipeline(
+        client=client,
         endpoint=endpoint,
-        version="v4",
-        schema=payload,
-        return_response=True,
-    )
-
-    wait_until_ticket_is_ready(client=client, ticket=ticket, verbose=verbose)
-
-    return client.call_endpoint(
-        "PileCore", "/get-task-result", version="v4", schema=ticket.json()
+        payload=payload,
+        verbose=verbose,
+        response_type="dict",
+        save_failed_payload=save_failed_payload,
+        failed_payload_filename=failed_payload_filename,
     )
 
 
@@ -254,6 +331,8 @@ def get_multi_cpt_api_report_tension(
     payload: dict,
     verbose: bool = False,
     standard: Literal["NEN9997-1", "CUR236"] = "NEN9997-1",
+    save_failed_payload: bool = False,
+    failed_payload_filename: str = "pilecore_multi_cpt_tension_report_debug_payload.json",
 ) -> bytes:
     """
     Wrapper around the PileCore endpoint "/uplift/[nen or cur]/multiple-cpts/report".
@@ -280,15 +359,12 @@ def get_multi_cpt_api_report_tension(
         endpoint = "/uplift/cur/multiple-cpts/report"
         payload.pop("construction_sequence", None)
 
-    ticket = client.call_endpoint(
-        "PileCore",
+    return get_task_result_pipeline(
+        client=client,
         endpoint=endpoint,
-        version="v4",
-        schema=payload,
-        return_response=True,
-    )
-    wait_until_ticket_is_ready(client=client, ticket=ticket, verbose=verbose)
-
-    return client.call_endpoint(
-        "PileCore", "/get-task-result", version="v4", schema=ticket.json()
+        payload=payload,
+        verbose=verbose,
+        response_type="bytes",
+        save_failed_payload=save_failed_payload,
+        failed_payload_filename=failed_payload_filename,
     )