cradle.sdk.types.tasks.common

  1from __future__ import annotations
  2
  3from abc import ABC
  4from enum import StrEnum
  5from typing import Literal
  6
  7from pydantic import BaseModel, Field, JsonValue
  8
  9
 10class TaskState(StrEnum):
 11    INIT = "INIT"
 12    PREPARING_INPUTS = "PREPARING_INPUTS"
 13    LAUNCHING = "LAUNCHING"
 14    EXECUTING = "EXECUTING"
 15    LOADING_RESULTS = "LOADING_RESULTS"
 16    CANCELLING = "CANCELLING"
 17    RECOVERING = "RECOVERING"
 18    FAILED = "FAILED"
 19    COMPLETED = "COMPLETED"
 20    CANCELLED = "CANCELLED"
 21
 22
 23class ArtifactParam(BaseModel):
 24    artifact_id: int
 25
 26
 27class CustomPredictorParam(BaseModel):
 28    """Reference to a custom predictor."""
 29
 30    name: str = Field(
 31        description="Specifies the name of the custom predictor to be used. Must match a custom predictor added to the workspace.",
 32        min_length=1,
 33        max_length=64,
 34    )
 35    version_id: int = Field(description="Version of the predictor to use. Must be a version in `ACTIVE` state.")
 36    parameters: dict[str, bool | int | float | str] = Field(
 37        default_factory=dict,
 38        description="Key-value map of predictor-specific parameters to pass to the predictor. For a given pair, `(key, value)`, `key` should correspond to one of the input names of the custom predictor and `value` should be of the appropriate type for that input. See also the `inputs` field in the documentation for creating custom predictors.",
 39    )
 40
 41
 42class ParametersBase(BaseModel, ABC):
 43    """Class that marks top-level task inputs. Used for static checks."""
 44
 45
 46class QueryableTableSource(BaseModel, ABC):
 47    """A table source that can produce a SQL query for BigQuery execution."""
 48
 49
 50class ResultBase(BaseModel, ABC):
 51    """Class that marks top-level task results. Used for static checks."""
 52
 53
 54class TableOutputBase(BaseModel):
 55    table: str = Field(description="The reference to the table the result was written to.")
 56
 57
 58class TableSourceLiteral(BaseModel):
 59    """Inline literal table data with full schema support (including nested structs).
 60
 61    Rows are validated against the expected table schema. Maximum 100 rows.
 62    """
 63
 64    kind: Literal["LITERAL"] = Field(default="LITERAL")
 65    rows: list[dict[str, JsonValue]] = Field(max_length=100)
 66
 67
 68class TableSourceSnapshot(BaseModel):
 69    kind: Literal["SNAPSHOT"] = Field(default="SNAPSHOT")
 70    id: int
 71
 72
 73class TableSourceQuery(QueryableTableSource):
 74    kind: Literal["QUERY"] = Field(default="QUERY")
 75    query: str = Field(max_length=100000)
 76
 77
 78class TableSourceTable(QueryableTableSource):
 79    kind: Literal["TABLE"] = Field(default="TABLE")
 80    table: str
 81
 82
 83class TableSourceTaskResult(QueryableTableSource):
 84    """A reference to the output of a previous task, optionally transformed by a SQL query.
 85
 86    Without ``query``: produces
 87    ``SELECT * EXCEPT(cradle, task_id) FROM <table> WHERE task_id IN (...)``.
 88
 89    With ``query``: the base result is injected as a CTE named ``__source__``, and the
 90    query is applied on top::
 91
 92        WITH __source__ AS (
 93            SELECT * EXCEPT(cradle, task_id) FROM <table> WHERE task_id IN (...)
 94        )
 95        <your query referencing __source__>
 96
 97    The query must reference ``__source__`` and must not reference ``cradle.results.*``
 98    tables (use a separate `TableSourceTaskResult` to declare that provenance link).
 99    """
100
101    kind: Literal["TASK_RESULT"] = Field(default="TASK_RESULT")
102    table: str
103    task_id: int | list[int]
104    query: str | None = Field(
105        default=None,
106        description="Optional SQL query applied on top of the task result. The base result is available as a CTE named `__source__`. Must reference `__source__` and must not reference `cradle.results.*` tables.",
107        max_length=100000,
108    )
109
110
111class TableResult(TableOutputBase):
112    pass