cradle.sdk.utils

 1import json
 2import time
 3from pathlib import Path
 4
 5from cradle.sdk.types.data import FileUploadResponse
 6from cradle.sdk.types.task import TaskResponse
 7
 8from .client import DataLoadClient, TaskClient
 9
10
11def upload_blob_files(
12    client: DataLoadClient, load_id: int, files: list[Path], table_reference: str | None = None
13) -> list[FileUploadResponse]:
14    responses = []
15    for f in files:
16        response = client.upload_file(load_id=load_id, file=f, table_reference=table_reference)
17        print(f"Uploaded {f}, id={response.id}")
18        responses.append(response)
19    return responses
20
21
22def wait_for_task(client: TaskClient, task_id: int, timeout: float = 60) -> TaskResponse:
23    start = time.monotonic()
24    while time.monotonic() - start < timeout:
25        task = client.get(task_id)
26        print(f"\r{task.state} @ {task.updated_at}", end="", flush=True)
27        if task.state == "COMPLETED":
28            if task.result is None:
29                raise ValueError("Task result is None")
30            print()
31            print(json.dumps(task.result.model_dump(mode="json"), indent=2))
32            return task
33        if task.state == "FAILED":
34            print()
35            print(task.errors)
36            return task
37        if task.state == "CANCELLED":
38            return task
39        time.sleep(1)
40    return client.get(task_id)
def upload_blob_files( client: cradle.sdk.client.DataLoadClient, load_id: int, files: list[pathlib.Path], table_reference: str | None = None) -> list[cradle.sdk.types.data.FileUploadResponse]:
12def upload_blob_files(
13    client: DataLoadClient, load_id: int, files: list[Path], table_reference: str | None = None
14) -> list[FileUploadResponse]:
15    responses = []
16    for f in files:
17        response = client.upload_file(load_id=load_id, file=f, table_reference=table_reference)
18        print(f"Uploaded {f}, id={response.id}")
19        responses.append(response)
20    return responses
def wait_for_task( client: cradle.sdk.client.TaskClient, task_id: int, timeout: float = 60) -> cradle.sdk.types.task.TaskResponse:
23def wait_for_task(client: TaskClient, task_id: int, timeout: float = 60) -> TaskResponse:
24    start = time.monotonic()
25    while time.monotonic() - start < timeout:
26        task = client.get(task_id)
27        print(f"\r{task.state} @ {task.updated_at}", end="", flush=True)
28        if task.state == "COMPLETED":
29            if task.result is None:
30                raise ValueError("Task result is None")
31            print()
32            print(json.dumps(task.result.model_dump(mode="json"), indent=2))
33            return task
34        if task.state == "FAILED":
35            print()
36            print(task.errors)
37            return task
38        if task.state == "CANCELLED":
39            return task
40        time.sleep(1)
41    return client.get(task_id)