cradle-sdk 0.1.1__py3-none-any.whl

cradle/sdk/exceptions.py

@@ -0,0 +1,11 @@
+class ClientError(Exception):
+    def __init__(self, status_code: int, message: str, errors: list[str]):
+        self.status_code = status_code
+        self.message = message
+
+        message = f"API Error {status_code}: {message}"
+        if len(errors) > 0:
+            error_list = "\n".join(f"\t{error}" for error in errors)
+            message = f"{message}\n{error_list}"
+
+        super().__init__(message)

cradle/sdk/types/assembly.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel, Field
+
+
+class AssemblyFeature(BaseModel):
+    ranges: dict[str, list[tuple[int, int]]] = Field(description="The ranges of the feature, keyed by polymer name")
+    metadata: dict[str, str | int | float | bool] | None = Field(
+        default=None, description="Optional metadata for the feature"
+    )

cradle/sdk/types/common.py

@@ -0,0 +1,62 @@
+from datetime import datetime
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+
+class ListOptions(BaseModel):
+    limit: int = Field(default=100, description="Maximum number of items to return")
+    cursor: str | None = Field(default=None, description="Continuation token returned by the previous page's response")
+    order: Literal["ASC", "DESC"] = Field(
+        default="ASC", description="Whether to order the result ascending or descending"
+    )
+
+
+class AbstractResourceResponse(BaseModel):
+    id: int = Field(description="The unique identifier of the resource")
+    created_at: datetime = Field(description="The time at which the resource was created")
+    created_by: str = Field(description="The ID of the user who created the resource")
+
+
+class ResourceResponse(AbstractResourceResponse):
+    updated_at: datetime = Field(description="The time at which the resource was last updated")
+    updated_by: str = Field(description="The ID of the user who last updated the resource")
+
+
+class ErrorResponseMixin(BaseModel):
+    """Mixin for resources that can return errors."""
+
+    errors: list[str] = Field(default_factory=list, description="List of error messages associated with the resource.")
+
+
+class BaseListResponse(BaseModel):
+    items: list[AbstractResourceResponse] = Field(description="The listed resources.")
+    cursor: str | None = Field(
+        description="Cursor with which the next page of entries can be retrieved. If no cursor is returned, no further pages exist."
+    )
+
+
+class ArchivableResourceMixin(BaseModel):
+    archived_at: datetime | None = Field(default=None, description="The time at which the resource was archived")
+    archived_by: str | None = Field(default=None, description="The ID of the user who archived the resource")
+
+
+class ImmutableResourceResponse(AbstractResourceResponse):
+    pass
+
+
+class ContextWorkspace(BaseModel):
+    kind: Literal["WORKSPACE"] = Field(default="WORKSPACE")
+
+
+class ContextProject(BaseModel):
+    kind: Literal["PROJECT"] = Field(default="PROJECT")
+    project_id: int
+
+
+class ContextRound(BaseModel):
+    kind: Literal["ROUND"] = Field(default="ROUND")
+    round_id: int
+
+
+Context = Annotated[ContextWorkspace | ContextProject | ContextRound, Field(discriminator="kind")]

cradle/sdk/types/data.py

@@ -0,0 +1,353 @@
+from enum import StrEnum
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from .common import (
+    ArchivableResourceMixin,
+    BaseListResponse,
+    Context,
+    ContextProject,
+    ContextRound,
+    ErrorResponseMixin,
+    ImmutableResourceResponse,
+    ResourceResponse,
+)
+
+
+class QueryDataRequest(BaseModel):
+    query: str = Field(description="The SQL query to execute.")
+    arrow_compression: Literal["zstd", "lz4"] | None = Field(
+        default=None,
+        description="The compression algorithm to use for the Arrow response. Supported values are `zstd`, `lz4`, and `None`(meaning no compression).",
+    )
+    project_id: int | None = Field(
+        default=None, description="If set, the query will only run over data associated with the given `project_id`."
+    )
+    version_id: int | None = Field(
+        default=None,
+        description="If set, will run the query against the given version of the data according to the changelog.",
+    )
+
+
+class DataActionKind(StrEnum):
+    LOAD = "LOAD"
+    UNDO_LOAD = "UNDO_LOAD"
+    TABLE_CREATE = "TABLE_CREATE"
+    TABLE_UPDATE = "TABLE_UPDATE"
+    TABLE_ARCHIVE = "TABLE_ARCHIVE"
+
+
+class DataActionLoad(BaseModel):
+    kind: Literal[DataActionKind.LOAD] = Field(default=DataActionKind.LOAD)
+    load_id: int = Field(description="The ID of the data load from which data was added to the dataset.")
+
+
+class DataActionUndoLoad(BaseModel):
+    kind: Literal[DataActionKind.UNDO_LOAD] = Field(default=DataActionKind.UNDO_LOAD)
+    load_id: int = Field(description="ID of the data load of which data was removed from the dataset.")
+
+
+class DataActionTableCreate(BaseModel):
+    kind: Literal[DataActionKind.TABLE_CREATE] = Field(default=DataActionKind.TABLE_CREATE)
+    table_id: int = Field(description="The ID of the table that was created.")
+    table_version_id: int = Field(description="The ID of the table version that was created.")
+
+
+class DataActionTableUpdate(BaseModel):
+    kind: Literal[DataActionKind.TABLE_UPDATE] = Field(default=DataActionKind.TABLE_UPDATE)
+    table_id: int = Field(description="The ID of the table that was updated.")
+    table_version_id: int = Field(description="The ID of the new table version.")
+
+
+class DataActionTableArchive(BaseModel):
+    kind: Literal[DataActionKind.TABLE_ARCHIVE] = Field(default=DataActionKind.TABLE_ARCHIVE)
+    table_id: int = Field(description="The ID of the table that was archived.")
+
+
+class DataVersionResponse(ImmutableResourceResponse):
+    action: (
+        DataActionLoad
+        | DataActionUndoLoad
+        | Annotated[DataActionTableCreate | DataActionTableUpdate | DataActionTableArchive, Field(discriminator="kind")]
+    ) = Field(discriminator="kind")
+
+
+class ListDataVersionResponse(BaseListResponse):
+    items: list[DataVersionResponse]
+
+
+class ArtifactProducerKind(StrEnum):
+    TASK = "TASK"
+    USER = "USER"
+
+
+class TaskProducer(BaseModel):
+    kind: Literal[ArtifactProducerKind.TASK] = Field(default=ArtifactProducerKind.TASK)
+    task_id: int
+
+
+class UserProducer(BaseModel):
+    kind: Literal[ArtifactProducerKind.USER] = Field(default=ArtifactProducerKind.USER)
+
+
+class ArtifactResponse(ImmutableResourceResponse):
+    context: "Context" = Field(discriminator="kind")
+    type: Literal[
+        "protein_structure/v1",
+        "scorer/v1",
+        "report/v1",
+        "sampler/v1",
+        "biostats/v1",
+        "scorer.monomer/v2",
+        "scorer.vhvl/v2",
+        "sampler.monomer/v2",
+        "sampler.vhvl/v2",
+        "analysis_data.train/v1",
+    ]
+    producer: TaskProducer | UserProducer = Field(discriminator="kind")
+
+
+class ListArtifactResponse(BaseListResponse):
+    items: list[ArtifactResponse]
+
+
+class TableKind(StrEnum):
+    TABLE = "TABLE"
+    VIEW = "VIEW"
+
+
+class ColumnMetadata(BaseModel):
+    description: str | None = Field(default=None)
+    deprecated: bool = Field(default=False)
+
+
+class ColumnBase(BaseModel):
+    name: str
+    metadata: ColumnMetadata | None = Field(default=None)
+
+
+class TypeNames(StrEnum):
+    BOOL = "BOOL"
+    INT64 = "INT64"
+    FLOAT64 = "FLOAT64"
+    STRING = "STRING"
+    STRUCT = "STRUCT"
+    ARRAY = "ARRAY"
+
+
+class PrimitiveType(BaseModel):
+    type: Literal[TypeNames.BOOL, TypeNames.INT64, TypeNames.FLOAT64, TypeNames.STRING]
+
+
+class PrimitiveColumn(ColumnBase, PrimitiveType):
+    """A column representing a primitive type."""
+
+    nullable: bool = Field(default=True)
+
+
+class ArrayType(BaseModel):
+    type: Literal[TypeNames.ARRAY] = Field(default=TypeNames.ARRAY)
+    item_type: "ColumnType"
+
+
+class ArrayColumn(ColumnBase, ArrayType):
+    """A column representing an array of elements.
+
+    Unlike the other column types it is not nullable but the array can be empty.
+    This matches BigQuery's behavior and reduces ambiguity.
+    """
+
+
+class StructType(BaseModel):
+    type: Literal[TypeNames.STRUCT] = Field(default=TypeNames.STRUCT)
+    columns: list["Column"]
+
+
+class StructColumn(ColumnBase, StructType):
+    """A column representing a struct of fields."""
+
+    nullable: bool = Field(default=True)
+
+
+class TableResponseBase(ResourceResponse, ArchivableResourceMixin):
+    """Base class for table responses."""
+
+    kind: TableKind
+    reference: str = Field(
+        description="""A unique name to use for referring to the table, such as `temperature_assay`. User-created
+tables should always use this simple `table_name` format. Cradle-created table references have the form
+`catalog.schema.table_name`, for example `cradle.results.engineered_sequences`.
+
+References must be unique per workspace. One exception to the rule is that once a table is archived the
+refence is available to be used again. References are tied to a data version, i.e. when running a query
+at a historic version, a reference will point to the table that held the reference at this data version."""
+    )
+    description: str | None
+    columns: list["Column"] = Field(
+        description="The columns of the table. For views, the columns and their types are inferred from the query at view creation or update time."
+    )
+    version_id: int = Field(
+        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."
+    )
+
+
+class BaseTableResponse(TableResponseBase):
+    kind: Literal[TableKind.TABLE] = Field(default=TableKind.TABLE)
+
+
+class ViewTableResponse(TableResponseBase):
+    kind: Literal[TableKind.VIEW] = Field(default=TableKind.VIEW)
+    query: str = Field(description="The SQL query that defines the view.")
+
+
+class ListTableResponse(BaseListResponse):
+    items: list["TableResponse"]
+
+
+class TableCreateBase(BaseModel):
+    kind: TableKind
+    reference: str = Field(
+        description="""A unique name to use for referring to the table, such as `temperature_assay`. User-created
+tables should always use this simple `table_name` format. Cradle-created table references have the form
+`catalog.schema.table_name`, for example `cradle.results.engineered_sequences`.
+
+References must be unique per workspace. One exception to the rule is that once a table is archived the
+refence is available to be used again. References are tied to a data version, i.e. when running a query
+at a historic version, a reference will point to the table that held the reference at this data version."""
+    )
+    description: str | None = Field(default=None, description="A human-friendly description of the table's contents.")
+
+
+class BaseTableCreate(TableCreateBase):
+    kind: Literal[TableKind.TABLE] = Field(
+        default=TableKind.TABLE,
+        description="""Whether to create a *table* or a *view*. Tables are created with the purpose of uploading data
+        to them, while views are created as queries on top of existing tables or views. Think of tables as physical
+        entities that store data, while views are virtual entities that represent a query result.
+        """,
+    )
+    columns: list["Column"] = Field(description="Columns of the table")
+
+
+class ViewTableCreate(TableCreateBase):
+    kind: Literal[TableKind.VIEW] = Field(
+        default=TableKind.VIEW,
+        description="""Whether to create a *table* or a *view*. Tables are created with the purpose of uploading data
+        to them, while views are created as queries on top of existing tables or views. Think of tables as physical
+        entities that store data, while views are virtual entities that represent a query result.
+        """,
+    )
+    query: str = Field(
+        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."
+    )
+    columns: list["Column"] | None = Field(
+        default=None,
+        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.",
+    )
+
+
+class TableRename(BaseModel):
+    new_reference: str = Field(
+        description="""A unique name to use for referring to the table, such as `temperature_assay`. User-created
+tables should always use this simple `table_name` format. Cradle-created table references have the form
+`catalog.schema.table_name`, for example `cradle.results.engineered_sequences`.
+
+References must be unique per workspace. One exception to the rule is that once a table is archived the
+refence is available to be used again. References are tied to a data version, i.e. when running a query
+at a historic version, a reference will point to the table that held the reference at this data version."""
+    )
+
+
+class _TableUpdateBase(BaseModel):
+    description: str | None = Field(default=None, description="Description of the table.")
+
+
+class BaseTableUpdate(_TableUpdateBase):
+    kind: Literal[TableKind.TABLE] = Field(default=TableKind.TABLE)
+    columns: list["Column"] | None = Field(default=None, description="Columns of the table.")
+
+
+class ViewTableUpdate(_TableUpdateBase):
+    kind: Literal[TableKind.VIEW] = Field(default=TableKind.VIEW)
+    query: str | None = Field(
+        default=None,
+        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.",
+    )
+    columns: list["Column"] | None = Field(
+        default=None,
+        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.",
+    )
+
+
+class TableArchive(BaseModel):
+    pass
+
+
+class FileFormat(StrEnum):
+    CSV = "CSV"
+    PARQUET = "PARQUET"
+    NDJSON = "NDJSON"
+
+
+class TableLoadConfig(BaseModel):
+    columns: list["Column"] = Field(
+        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."
+    )
+    format: FileFormat
+
+
+class DataLoadCreate(BaseModel):
+    context: ContextProject | ContextRound = Field(
+        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)."
+    )
+    tables: dict[str, TableLoadConfig] = Field(
+        default_factory=dict,
+        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.",
+    )
+
+
+class TableLoadConfigResponse(TableLoadConfig):
+    table_id: int | None = Field(
+        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."
+    )
+
+
+class DataLoadState(StrEnum):
+    PENDING = "PENDING"
+    LOADING = "LOADING"
+    DELETING = "DELETING"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+
+
+class FileUploadResponse(ImmutableResourceResponse):
+    filepath: str
+    size: int
+    table_reference: str | None
+    source_file_id: int | None
+    description: str | None
+
+
+class DataLoadResponse(ResourceResponse, ErrorResponseMixin):
+    context: "Context" = Field(description="The context the loaded data is be associated with", discriminator="kind")
+    tables: dict[str, TableLoadConfigResponse] = Field(
+        default_factory=dict, description="The tables for which data will be loaded."
+    )
+    state: DataLoadState
+    is_active: bool
+    files: list[FileUploadResponse] = Field(description="Files that have been uploaded to this data load.")
+
+
+class AddTableRequest(BaseModel):
+    table_ref: str
+    config: TableLoadConfig
+
+
+class ListDataLoadResponse(BaseListResponse):
+    items: list[DataLoadResponse]
+
+
+Column = Annotated[PrimitiveColumn | StructColumn | ArrayColumn, Field(discriminator="type")]
+ColumnType = Annotated[PrimitiveType | StructType | ArrayType, Field(discriminator="type")]
+TableResponse = Annotated[BaseTableResponse | ViewTableResponse, Field(discriminator="kind")]