cradle.sdk.types.data

  1from __future__ import annotations
  2
  3from enum import StrEnum
  4from typing import Annotated, Any, Literal, TypeAlias
  5
  6from pydantic import BaseModel, Field
  7
  8from .common import (
  9    ArchivableResourceMixin,
 10    BaseListResponse,
 11    ContextProject,
 12    ContextRound,
 13    ContextWorkspace,
 14    ErrorResponseMixin,
 15    ImmutableResourceResponse,
 16    ResourceResponse,
 17)
 18
 19
 20class ArtifactProducerKind(StrEnum):
 21    TASK = "TASK"
 22    USER = "USER"
 23
 24
 25class DataActionKind(StrEnum):
 26    LOAD = "LOAD"
 27    UNDO_LOAD = "UNDO_LOAD"
 28    TABLE_CREATE = "TABLE_CREATE"
 29    TABLE_UPDATE = "TABLE_UPDATE"
 30    TABLE_ARCHIVE = "TABLE_ARCHIVE"
 31
 32
 33class DataLoadState(StrEnum):
 34    PENDING = "PENDING"
 35    LOADING = "LOADING"
 36    DELETING = "DELETING"
 37    COMPLETED = "COMPLETED"
 38    FAILED = "FAILED"
 39
 40
 41class FileFormat(StrEnum):
 42    CSV = "CSV"
 43    PARQUET = "PARQUET"
 44    NDJSON = "NDJSON"
 45
 46
 47class TableKind(StrEnum):
 48    TABLE = "TABLE"
 49    VIEW = "VIEW"
 50
 51
 52class TypeNames(StrEnum):
 53    BOOL = "BOOL"
 54    INT64 = "INT64"
 55    FLOAT64 = "FLOAT64"
 56    STRING = "STRING"
 57    STRUCT = "STRUCT"
 58    ARRAY = "ARRAY"
 59
 60
 61class AddTableRequest(BaseModel):
 62    table_ref: str
 63    config: TableLoadConfig
 64
 65
 66class ColumnBase(BaseModel):
 67    name: str = Field(min_length=1, max_length=64)
 68    metadata: ColumnMetadata | None = Field(default=None)
 69
 70
 71class ArrayType(BaseModel):
 72    type: Literal[TypeNames.ARRAY] = Field(default=TypeNames.ARRAY)
 73    item_type: PrimitiveType | StructType | ArrayType = Field(discriminator="type")
 74
 75
 76class ArtifactResponse(ImmutableResourceResponse):
 77    context: ContextWorkspace | ContextProject | ContextRound = Field(discriminator="kind")
 78    type: str
 79    data: dict[str, Any]
 80    producer: TaskProducer | UserProducer = Field(discriminator="kind")
 81
 82
 83class TableCreateBase(BaseModel):
 84    kind: TableKind
 85    reference: str = Field(
 86        description="A unique name to use for referring to the table, such as `temperature_assay`. User-created\ntables should always use this simple `table_name` format. Cradle-created table references have the form\n`catalog.schema.table_name`, for example `cradle.results.engineered_sequences`.\n\nReferences must be unique per workspace. One exception to the rule is that once a table is archived the\nreference is available to be used again. References are tied to a data version, i.e. when running a query\nat a historic version, a reference will point to the table that held the reference at this data version."
 87    )
 88    description: str | None = Field(
 89        default=None, description="A human-friendly description of the table's contents.", min_length=1, max_length=1024
 90    )
 91
 92
 93class TableResponseBase(ResourceResponse, ArchivableResourceMixin):
 94    """Base class for table responses."""
 95
 96    kind: TableKind
 97    reference: str = Field(
 98        description="A unique name to use for referring to the table, such as `temperature_assay`. User-created\ntables should always use this simple `table_name` format. Cradle-created table references have the form\n`catalog.schema.table_name`, for example `cradle.results.engineered_sequences`.\n\nReferences must be unique per workspace. One exception to the rule is that once a table is archived the\nreference is available to be used again. References are tied to a data version, i.e. when running a query\nat a historic version, a reference will point to the table that held the reference at this data version."
 99    )
100    description: str | None
101    columns: list[Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]] = Field(
102        description="The columns of the table. For views, the columns and their types are inferred from the query at view creation or update time."
103    )
104    version_id: int = Field(
105        description="The version ID of the table state reflected in the response. Changes to a view's query or a base table's columns cause a version change as well as renames cause a version change.",
106        deprecated="Table versioning will no longer be tracked in future platform versions",
107    )
108
109
110class _TableUpdateBase(BaseModel):
111    description: str | None = Field(
112        default=None, description="Description of the table.", min_length=1, max_length=1024
113    )
114
115
116class ColumnMetadata(BaseModel):
117    description: str | None = Field(default=None, min_length=1, max_length=1024)
118    deprecated: bool = Field(default=False)
119
120
121class DataActionLoad(BaseModel):
122    kind: Literal[DataActionKind.LOAD] = Field(default=DataActionKind.LOAD)
123    load_id: int = Field(description="The ID of the data load from which data was added to the dataset.")
124
125
126class DataActionTableArchive(BaseModel):
127    kind: Literal[DataActionKind.TABLE_ARCHIVE] = Field(default=DataActionKind.TABLE_ARCHIVE)
128    table_id: int = Field(description="The ID of the table that was archived.")
129
130
131class DataActionTableCreate(BaseModel):
132    kind: Literal[DataActionKind.TABLE_CREATE] = Field(default=DataActionKind.TABLE_CREATE)
133    table_id: int = Field(description="The ID of the table that was created.")
134    table_version_id: int = Field(description="The ID of the table version that was created.")
135
136
137class DataActionTableUpdate(BaseModel):
138    kind: Literal[DataActionKind.TABLE_UPDATE] = Field(default=DataActionKind.TABLE_UPDATE)
139    table_id: int = Field(description="The ID of the table that was updated.")
140    table_version_id: int = Field(description="The ID of the new table version.")
141
142
143class DataActionUndoLoad(BaseModel):
144    kind: Literal[DataActionKind.UNDO_LOAD] = Field(default=DataActionKind.UNDO_LOAD)
145    load_id: int = Field(description="ID of the data load of which data was removed from the dataset.")
146
147
148class DataLoadCreate(BaseModel):
149    context: ContextProject | ContextRound = Field(
150        description="The context the loaded data will be associated with. The data will either live within a Project (visible to all rounds) or within a specific Round of a project (visible only to that round)."
151    )
152    tables: dict[str, TableLoadConfig] = Field(
153        default_factory=dict,
154        description="A map of table references to table schemas. The map keys define the tables for which data is uploaded. The map values define the schema and file format for a given table.",
155    )
156
157
158class DataLoadResponse(ResourceResponse, ErrorResponseMixin):
159    context: ContextWorkspace | ContextProject | ContextRound = Field(
160        description="The context the loaded data is be associated with", discriminator="kind"
161    )
162    tables: dict[str, TableLoadConfigResponse] = Field(
163        default_factory=dict, description="The tables for which data will be loaded."
164    )
165    state: DataLoadState
166    is_active: bool
167    files: list[FileUploadResponse] = Field(description="Files that have been uploaded to this data load.")
168
169
170class DataVersionResponse(ImmutableResourceResponse):
171    action: (
172        DataActionLoad
173        | DataActionUndoLoad
174        | Annotated[DataActionTableCreate | DataActionTableUpdate | DataActionTableArchive, Field(discriminator="kind")]
175    ) = Field(discriminator="kind")
176
177
178class FileUploadResponse(ImmutableResourceResponse):
179    filepath: str
180    size: int
181    table_reference: str | None
182    source_file_id: int | None
183    description: str | None
184
185
186class ListArtifactResponse(BaseListResponse):
187    items: list[ArtifactResponse]
188
189
190class ListDataLoadResponse(BaseListResponse):
191    items: list[DataLoadResponse]
192
193
194class ListDataVersionResponse(BaseListResponse):
195    items: list[DataVersionResponse]
196
197
198class ListTableResponse(BaseListResponse):
199    items: list[Annotated[BaseTableResponse | ViewTableResponse, Field(discriminator="kind")]]
200
201
202class PrimitiveType(BaseModel):
203    type: Literal[TypeNames.BOOL, TypeNames.INT64, TypeNames.FLOAT64, TypeNames.STRING]
204
205
206class QueryDataRequest(BaseModel):
207    query: str = Field(description="The SQL query to execute.", max_length=100000)
208    arrow_compression: Literal["zstd", "lz4"] | None = Field(
209        default=None,
210        description="The compression algorithm to use for the Arrow response. Supported values are `zstd`, `lz4`, and `None`(meaning no compression).",
211    )
212    project_id: int | None = Field(
213        default=None, description="If set, the query will only run over data associated with the given `project_id`."
214    )
215    version_id: int | None = Field(
216        default=None,
217        description="If set, will run the query against the given version of the data according to the changelog.",
218        deprecated="Query specific data snapshots instead",
219    )
220
221
222class StructType(BaseModel):
223    type: Literal[TypeNames.STRUCT] = Field(default=TypeNames.STRUCT)
224    columns: list[Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]] = Field(
225        min_length=1
226    )
227
228
229class TableArchive(BaseModel):
230    pass
231
232
233class TableLoadConfig(BaseModel):
234    columns: list[Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]] = Field(
235        description="The column layout of the files. All files in the data load destined for this table must have this exact column layout. For CSV files, only the specified columns determine how data is loaded. The first row (header) in each uploaded file is skipped. For Parquet files, columns will be matched by name and are thus not order-sensitive. Therefore it is recommended to use Parquet files for more robust data loading.",
236        min_length=1,
237    )
238    format: FileFormat
239
240
241class TableRename(BaseModel):
242    new_reference: str = Field(
243        description="A unique name to use for referring to the table, such as `temperature_assay`. User-created\ntables should always use this simple `table_name` format. Cradle-created table references have the form\n`catalog.schema.table_name`, for example `cradle.results.engineered_sequences`.\n\nReferences must be unique per workspace. One exception to the rule is that once a table is archived the\nreference is available to be used again. References are tied to a data version, i.e. when running a query\nat a historic version, a reference will point to the table that held the reference at this data version."
244    )
245
246
247class TaskProducer(BaseModel):
248    kind: Literal[ArtifactProducerKind.TASK] = Field(default=ArtifactProducerKind.TASK)
249    task_id: int
250
251
252class UserProducer(BaseModel):
253    kind: Literal[ArtifactProducerKind.USER] = Field(default=ArtifactProducerKind.USER)
254
255
256class ArrayColumn(ColumnBase, ArrayType):
257    """A column representing an array of elements.
258
259    Unlike the other column types it is not nullable but the array can be empty.
260    This matches BigQuery's behavior and reduces ambiguity.
261    """
262
263
264class BaseTableCreate(TableCreateBase):
265    kind: Literal[TableKind.TABLE] = Field(
266        default=TableKind.TABLE,
267        description="Whether to create a *table* or a *view*. Tables are created with the purpose of uploading data\n        to them, while views are created as queries on top of existing tables or views. Think of tables as physical\n        entities that store data, while views are virtual entities that represent a query result.\n        ",
268    )
269    columns: list[Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]] = Field(
270        description="Columns of the table", min_length=1
271    )
272
273
274class ViewTableCreate(TableCreateBase):
275    kind: Literal[TableKind.VIEW] = Field(
276        default=TableKind.VIEW,
277        description="Whether to create a *table* or a *view*. Tables are created with the purpose of uploading data\n        to them, while views are created as queries on top of existing tables or views. Think of tables as physical\n        entities that store data, while views are virtual entities that represent a query result.\n        ",
278    )
279    query: str = Field(
280        description="SQL query for the view. It must be a single SELECT statement with a limited, but still very powerful, subset of SQL. Subqueries, common table expressions, joins, and conditional expressions with all commonly used SQL operators are supported. Contact your customer success representative if you need to use a specific SQL feature that is not supported. If you have a query use case that is not covered by the currently permitted SQL expressions, contact support to discuss potential extensions.",
281        min_length=1,
282        max_length=100000,
283    )
284    columns: (
285        Annotated[
286            list[Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]],
287            Field(min_length=1),
288        ]
289        | None
290    ) = Field(
291        default=None,
292        description="Columns of the view. The column schema must match the query's result schema. If unspecified, a schema is automatically inferred from the query and column metadata will be copied over from the existing schema if the column already existed.",
293    )
294
295
296class BaseTableResponse(TableResponseBase):
297    kind: Literal[TableKind.TABLE] = Field(default=TableKind.TABLE)
298
299
300class ViewTableResponse(TableResponseBase):
301    kind: Literal[TableKind.VIEW] = Field(default=TableKind.VIEW)
302    query: str = Field(description="The SQL query that defines the view.")
303
304
305class BaseTableUpdate(_TableUpdateBase):
306    kind: Literal[TableKind.TABLE] = Field(default=TableKind.TABLE)
307    columns: (
308        Annotated[
309            list[Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]],
310            Field(min_length=1),
311        ]
312        | None
313    ) = Field(default=None, description="Columns of the table.")
314
315
316class ViewTableUpdate(_TableUpdateBase):
317    kind: Literal[TableKind.VIEW] = Field(default=TableKind.VIEW)
318    query: str | None = Field(
319        default=None,
320        description="SQL query for the view. It must be a single SELECT statement and only a subset of SQL is allowed. If unspecified, no change is made to the view query.If you have a query use case that is not covered by the currently permitted SQL expressions, contact support to discuss potential extensions.",
321        min_length=1,
322        max_length=100000,
323    )
324    columns: (
325        Annotated[
326            list[Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]],
327            Field(min_length=1),
328        ]
329        | None
330    ) = Field(
331        default=None,
332        description="Columns of the view. The column schema must match the query's result schema. If unspecified, a schema is automatically inferred from the query and column metadata will be copied over from the existing schema if the column already existed.",
333    )
334
335
336class PrimitiveColumn(ColumnBase, PrimitiveType):
337    """A column representing a primitive type."""
338
339    nullable: bool = Field(default=True)
340
341
342class StructColumn(ColumnBase, StructType):
343    """A column representing a struct of fields."""
344
345    nullable: bool = Field(default=True)
346
347
348class TableLoadConfigResponse(TableLoadConfig):
349    table_id: int | None = Field(
350        description="The table the rows were loaded into. This is the table addressed by the table reference for this configuration at the time of finalization."
351    )
352
353
354Column: TypeAlias = Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]
355ColumnType: TypeAlias = Annotated[PrimitiveType | StructType | ArrayType, Field(discriminator="type")]
356TableResponse: TypeAlias = Annotated[BaseTableResponse | ViewTableResponse, Field(discriminator="kind")]