elementary-runtime 0.1.0rc1__py3-none-any.whl

elementary_runtime/protocol/base.py

@@ -0,0 +1,82 @@
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Annotated
+
+from pydantic import (
+    AfterValidator,
+    AwareDatetime,
+    BaseModel,
+    ConfigDict,
+    PlainSerializer,
+)
+from pydantic.types import StringConstraints
+
+from elementary_runtime.protocol.limits import (
+    MAX_DISPLAY_NAME_LENGTH,
+    MAX_ERROR_MESSAGE_LENGTH,
+    MAX_ID_LENGTH,
+    MAX_INSTANCE_ID_LENGTH,
+    MAX_INSTANCE_ID_PREFIX_LENGTH,
+    MAX_NON_EMPTY_STR_LENGTH,
+)
+
+INSTANCE_ID_PATTERN = r"^[A-Za-z0-9][A-Za-z0-9._:/-]*$"
+
+
+class ProtocolModel(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False)
+
+
+class ProtocolEnum(str, Enum):
+    def __str__(self) -> str:
+        return str(self.value)
+
+
+def to_utc(value: datetime) -> datetime:
+    return value.astimezone(timezone.utc)
+
+
+def utc_isoformat(value: datetime) -> str:
+    return value.isoformat().replace("+00:00", "Z")
+
+
+UtcDatetime = Annotated[
+    AwareDatetime,
+    AfterValidator(to_utc),
+    PlainSerializer(utc_isoformat, when_used="json"),
+]
+
+InstanceId = Annotated[
+    str,
+    StringConstraints(
+        min_length=1,
+        max_length=MAX_INSTANCE_ID_LENGTH,
+        pattern=INSTANCE_ID_PATTERN,
+    ),
+]
+
+InstanceIdPrefix = Annotated[
+    str,
+    StringConstraints(
+        min_length=1,
+        max_length=MAX_INSTANCE_ID_PREFIX_LENGTH,
+        pattern=INSTANCE_ID_PATTERN,
+    ),
+]
+
+IdStr = Annotated[str, StringConstraints(min_length=1, max_length=MAX_ID_LENGTH)]
+
+NonEmptyStr = Annotated[
+    str,
+    StringConstraints(min_length=1, max_length=MAX_NON_EMPTY_STR_LENGTH),
+]
+
+DisplayName = Annotated[
+    str,
+    StringConstraints(min_length=1, max_length=MAX_DISPLAY_NAME_LENGTH),
+]
+
+ErrorMessage = Annotated[
+    str,
+    StringConstraints(min_length=1, max_length=MAX_ERROR_MESSAGE_LENGTH),
+]

elementary_runtime/protocol/enums.py

@@ -0,0 +1,89 @@
+from elementary_runtime.protocol.base import ProtocolEnum
+
+
+class TaskType(ProtocolEnum):
+    CUSTOM_SQL_TEST_V1 = "custom_sql_test.v1"
+    RUN_SQL_V1 = "run_sql.v1"
+
+
+class ConnectionHealth(ProtocolEnum):
+    HEALTHY = "healthy"
+    UNHEALTHY = "unhealthy"
+    NOT_CHECKED = "not_checked"
+
+
+class ConnectionErrorCode(ProtocolEnum):
+    AUTHENTICATION_FAILED = "authentication_failed"
+    NETWORK_UNREACHABLE = "network_unreachable"
+    PERMISSION_DENIED = "permission_denied"
+    QUERY_TIMEOUT = "query_timeout"
+    QUERY_FAILED = "query_failed"
+
+
+class HeartbeatStatus(ProtocolEnum):
+    EXTENDED = "extended"
+    STALE = "stale"
+    COMPLETED = "completed"
+    NOT_FOUND = "not_found"
+
+
+class SubmittedResultStatus(ProtocolEnum):
+    ACCEPTED = "accepted"
+    ALREADY_COMPLETED = "already_completed"
+
+
+class ErrorCode(ProtocolEnum):
+    INVALID_REQUEST = "invalid_request"
+    INVALID_TOKEN = "invalid_token"
+    FORBIDDEN_CONNECTION = "forbidden_connection"
+    UNKNOWN_INSTANCE = "unknown_instance"
+    UNKNOWN_RESOURCE = "unknown_resource"
+    INCOMPATIBLE_VERSION = "incompatible_version"
+    INVALID_CONNECTION_STATE = "invalid_connection_state"
+    PAYLOAD_TOO_LARGE = "payload_too_large"
+    UNSUPPORTED_TASK_TYPE = "unsupported_task_type"
+    INVALID_RESULT = "invalid_result"
+    INVALID_TELEMETRY_EVENT = "invalid_telemetry_event"
+    RATE_LIMITED = "rate_limited"
+    PROCESSING_FAILED = "processing_failed"
+    SERVICE_UNAVAILABLE = "service_unavailable"
+
+
+class TelemetryEventType(ProtocolEnum):
+    RUNTIME_STARTED = "runtime_started"
+    RUNTIME_STOPPED = "runtime_stopped"
+    TASK_EXECUTION_SUMMARY = "task_execution_summary"
+    RUNTIME_EXCEPTION = "runtime_exception"
+
+
+class RuntimeStopReason(ProtocolEnum):
+    GRACEFUL = "graceful"
+    DRAIN_TIMEOUT = "drain_timeout"
+
+
+class TaskExecutionOutcome(ProtocolEnum):
+    COMPLETED = "completed"
+    QUERY_ERROR = "query_error"
+    TIMEOUT = "timeout"
+    PERMISSION_ERROR = "permission_error"
+
+
+class ExceptionComponent(ProtocolEnum):
+    COORDINATOR = "coordinator"
+    SQL_VALIDATOR = "sql_validator"
+    WAREHOUSE_CLIENT = "warehouse_client"
+    RESULT_ENCODER = "result_encoder"
+
+
+class ExceptionOperation(ProtocolEnum):
+    REGISTER = "register"
+    POLL = "poll"
+    HEARTBEAT = "heartbeat"
+    EXECUTE = "execute"
+    SUBMIT_RESULTS = "submit_results"
+    SEND_TELEMETRY = "send_telemetry"
+    SHUTDOWN = "shutdown"
+
+
+class TestSourceType(ProtocolEnum):
+    CLOUD_TEST = "cloud_test"

elementary_runtime/protocol/errors.py

@@ -0,0 +1,28 @@
+from elementary_runtime.protocol.base import (
+    ErrorMessage,
+    IdStr,
+    NonEmptyStr,
+    ProtocolModel,
+)
+from elementary_runtime.protocol.enums import ErrorCode
+
+
+class TaskErrorDetail(ProtocolModel):
+    message: ErrorMessage
+
+
+class ErrorDetails(ProtocolModel):
+    field: NonEmptyStr | None = None
+    task_id: IdStr | None = None
+
+
+class ErrorBody(ProtocolModel):
+    code: ErrorCode
+    message: ErrorMessage
+    request_id: IdStr
+    retryable: bool
+    details: ErrorDetails | None
+
+
+class ErrorResponse(ProtocolModel):
+    error: ErrorBody

elementary_runtime/protocol/http.py

@@ -0,0 +1,15 @@
+from typing import Final
+
+
+class Paths:
+    REGISTER_INSTANCE: Final = "/instances/register"
+    POLL_TASKS: Final = "/tasks/poll"
+    HEARTBEAT_TASKS: Final = "/tasks/heartbeat"
+    SUBMIT_RESULTS: Final = "/tasks/results"
+    TELEMETRY: Final = "/telemetry"
+
+
+class Headers:
+    INSTANCE_ID: Final = "X-Elementary-Runtime-Instance-ID"
+    REQUEST_ID: Final = "X-Request-ID"
+    RETRY_AFTER: Final = "Retry-After"

elementary_runtime/protocol/instances.py

@@ -0,0 +1,51 @@
+from typing import Annotated
+
+from pydantic import AfterValidator, Field, model_validator
+from typing_extensions import Self
+
+from elementary_runtime.common.validators import require_unique, require_unique_attr
+from elementary_runtime.protocol.base import (
+    DisplayName,
+    IdStr,
+    NonEmptyStr,
+    ProtocolModel,
+)
+from elementary_runtime.protocol.enums import (
+    ConnectionErrorCode,
+    ConnectionHealth,
+    TaskType,
+)
+from elementary_runtime.protocol.limits import MAX_CONNECTIONS_PER_REGISTER
+
+
+class ObservedWarehouseIdentity(ProtocolModel):
+    warehouse_type: NonEmptyStr
+    logical_identifier: NonEmptyStr | None = None
+    database: NonEmptyStr | None = None
+
+
+class ConnectionReport(ProtocolModel):
+    warehouse_connection_id: IdStr
+    health: ConnectionHealth
+    observed_warehouse_identity: ObservedWarehouseIdentity | None
+    error_code: ConnectionErrorCode | None
+
+    @model_validator(mode="after")
+    def error_code_matches_health(self) -> Self:
+        if self.health is ConnectionHealth.UNHEALTHY:
+            if self.error_code is None:
+                raise ValueError("error_code is required when health is unhealthy")
+        elif self.error_code is not None:
+            raise ValueError("error_code must be null when health is not unhealthy")
+        return self
+
+
+class RegisterInstanceRequest(ProtocolModel):
+    display_name: DisplayName | None
+    agent_version: NonEmptyStr
+    supported_task_types: Annotated[list[TaskType], AfterValidator(require_unique)]
+    connections: Annotated[
+        list[ConnectionReport],
+        Field(max_length=MAX_CONNECTIONS_PER_REGISTER),
+        AfterValidator(require_unique_attr("warehouse_connection_id")),
+    ]

elementary_runtime/protocol/limits.py

@@ -0,0 +1,18 @@
+MAX_POLL_TASKS = 256
+MAX_HEARTBEAT_TASKS = 256
+MAX_RESULTS_PER_BATCH = 64
+MAX_TELEMETRY_EVENTS = 256
+MAX_CONNECTIONS_PER_REGISTER = 256
+
+MAX_QUERY_TIMEOUT_SECONDS = 86_400
+MAX_SAMPLE_ROW_LIMIT = 100_000
+MAX_SAMPLE_ROWS = 100_000
+MAX_SAMPLE_COLUMNS = 1_000
+
+MAX_ERROR_MESSAGE_LENGTH = 4_096
+MAX_EXCEPTION_FRAMES = 50
+MAX_INSTANCE_ID_LENGTH = 255
+MAX_INSTANCE_ID_PREFIX_LENGTH = 128
+MAX_DISPLAY_NAME_LENGTH = 255
+MAX_ID_LENGTH = 255
+MAX_NON_EMPTY_STR_LENGTH = 255

elementary_runtime/protocol/tasks/__init__.py

@@ -0,0 +1,80 @@
+from elementary_runtime.protocol.errors import TaskErrorDetail
+from elementary_runtime.protocol.tasks.custom_sql_test import (
+    CloudTestSource,
+    CustomSqlTestResult,
+    CustomSqlTestResultPayload,
+    CustomSqlTestTask,
+    CustomSqlTestTaskPayload,
+)
+from elementary_runtime.protocol.tasks.envelope import (
+    TaskEnvelope,
+    TaskResultEnvelope,
+    WarehouseTaskEnvelope,
+)
+from elementary_runtime.protocol.tasks.heartbeat import (
+    HeartbeatTasksRequest,
+    HeartbeatTasksResponse,
+    TaskHeartbeat,
+    TaskHeartbeatStatus,
+)
+from elementary_runtime.protocol.tasks.poll import PollTasksRequest, PollTasksResponse
+from elementary_runtime.protocol.tasks.registry import (
+    TASK_MODELS,
+    Task,
+    TaskModels,
+    TaskResult,
+)
+from elementary_runtime.protocol.tasks.results import (
+    SubmitResultsRequest,
+    SubmitResultsResponse,
+    SubmittedResult,
+)
+from elementary_runtime.protocol.tasks.run_sql import (
+    RunSqlResult,
+    RunSqlResultPayload,
+    RunSqlTask,
+    RunSqlTaskPayload,
+)
+from elementary_runtime.protocol.tasks.sql_query import (
+    ResultSample,
+    SampleColumnType,
+    SampleHeader,
+    SqlExecutionStatus,
+    SqlQueryResultPayload,
+    SqlQueryTaskPayload,
+)
+
+__all__ = [
+    "TASK_MODELS",
+    "CloudTestSource",
+    "CustomSqlTestResult",
+    "CustomSqlTestResultPayload",
+    "CustomSqlTestTask",
+    "CustomSqlTestTaskPayload",
+    "HeartbeatTasksRequest",
+    "HeartbeatTasksResponse",
+    "PollTasksRequest",
+    "PollTasksResponse",
+    "ResultSample",
+    "RunSqlResult",
+    "RunSqlResultPayload",
+    "RunSqlTask",
+    "RunSqlTaskPayload",
+    "SampleColumnType",
+    "SampleHeader",
+    "SqlExecutionStatus",
+    "SqlQueryResultPayload",
+    "SqlQueryTaskPayload",
+    "SubmitResultsRequest",
+    "SubmitResultsResponse",
+    "SubmittedResult",
+    "Task",
+    "TaskEnvelope",
+    "TaskErrorDetail",
+    "TaskHeartbeat",
+    "TaskHeartbeatStatus",
+    "TaskModels",
+    "TaskResult",
+    "TaskResultEnvelope",
+    "WarehouseTaskEnvelope",
+]

elementary_runtime/protocol/tasks/custom_sql_test.py

@@ -0,0 +1,35 @@
+from typing import Literal
+
+from elementary_runtime.protocol.base import IdStr, ProtocolModel
+from elementary_runtime.protocol.enums import TaskType, TestSourceType
+from elementary_runtime.protocol.tasks.envelope import (
+    TaskResultEnvelope,
+    WarehouseTaskEnvelope,
+)
+from elementary_runtime.protocol.tasks.sql_query import (
+    SqlQueryResultPayload,
+    SqlQueryTaskPayload,
+)
+
+
+class CloudTestSource(ProtocolModel):
+    type: Literal[TestSourceType.CLOUD_TEST]
+    id: IdStr
+
+
+class CustomSqlTestTaskPayload(SqlQueryTaskPayload):
+    source: CloudTestSource
+
+
+class CustomSqlTestTask(WarehouseTaskEnvelope):
+    task_type: Literal[TaskType.CUSTOM_SQL_TEST_V1]
+    payload: CustomSqlTestTaskPayload
+
+
+class CustomSqlTestResultPayload(SqlQueryResultPayload):
+    pass
+
+
+class CustomSqlTestResult(TaskResultEnvelope):
+    task_type: Literal[TaskType.CUSTOM_SQL_TEST_V1]
+    payload: CustomSqlTestResultPayload

elementary_runtime/protocol/tasks/envelope.py

@@ -0,0 +1,20 @@
+from typing import Annotated
+
+from pydantic import Field
+
+from elementary_runtime.protocol.base import IdStr, ProtocolModel
+
+
+class TaskEnvelope(ProtocolModel):
+    task_id: IdStr
+    delivery_id: IdStr
+
+
+class WarehouseTaskEnvelope(TaskEnvelope):
+    warehouse_connection_id: IdStr
+
+
+class TaskResultEnvelope(ProtocolModel):
+    task_id: IdStr
+    delivery_id: IdStr
+    duration_seconds: Annotated[float, Field(ge=0.0)]

elementary_runtime/protocol/tasks/heartbeat.py

@@ -0,0 +1,34 @@
+from typing import Annotated
+
+from pydantic import AfterValidator, Field
+
+from elementary_runtime.common.validators import require_unique_attr
+from elementary_runtime.protocol.base import IdStr, ProtocolModel
+from elementary_runtime.protocol.enums import HeartbeatStatus
+from elementary_runtime.protocol.limits import MAX_HEARTBEAT_TASKS
+
+
+class TaskHeartbeat(ProtocolModel):
+    task_id: IdStr
+    delivery_id: IdStr
+
+
+class HeartbeatTasksRequest(ProtocolModel):
+    tasks: Annotated[
+        list[TaskHeartbeat],
+        Field(max_length=MAX_HEARTBEAT_TASKS),
+        AfterValidator(require_unique_attr("task_id")),
+    ]
+
+
+class TaskHeartbeatStatus(ProtocolModel):
+    task_id: IdStr
+    status: HeartbeatStatus
+
+
+class HeartbeatTasksResponse(ProtocolModel):
+    tasks: Annotated[
+        list[TaskHeartbeatStatus],
+        Field(max_length=MAX_HEARTBEAT_TASKS),
+        AfterValidator(require_unique_attr("task_id")),
+    ]

elementary_runtime/protocol/tasks/poll.py

@@ -0,0 +1,20 @@
+from typing import Annotated
+
+from pydantic import AfterValidator, Field
+
+from elementary_runtime.common.validators import require_unique_attr
+from elementary_runtime.protocol.base import ProtocolModel
+from elementary_runtime.protocol.limits import MAX_POLL_TASKS
+from elementary_runtime.protocol.tasks.registry import Task
+
+
+class PollTasksRequest(ProtocolModel):
+    available_slots: Annotated[int, Field(ge=0, le=MAX_POLL_TASKS)]
+
+
+class PollTasksResponse(ProtocolModel):
+    tasks: Annotated[
+        list[Task],
+        Field(max_length=MAX_POLL_TASKS),
+        AfterValidator(require_unique_attr("task_id")),
+    ]

elementary_runtime/protocol/tasks/registry.py

@@ -0,0 +1,35 @@
+from collections.abc import Mapping
+from typing import Annotated, NamedTuple
+
+from pydantic import Field
+
+from elementary_runtime.protocol.enums import TaskType
+from elementary_runtime.protocol.tasks.custom_sql_test import (
+    CustomSqlTestResult,
+    CustomSqlTestTask,
+)
+from elementary_runtime.protocol.tasks.envelope import TaskEnvelope, TaskResultEnvelope
+from elementary_runtime.protocol.tasks.run_sql import RunSqlResult, RunSqlTask
+
+Task = Annotated[CustomSqlTestTask | RunSqlTask, Field(discriminator="task_type")]
+TaskResult = Annotated[
+    CustomSqlTestResult | RunSqlResult,
+    Field(discriminator="task_type"),
+]
+
+
+class TaskModels(NamedTuple):
+    task: type[TaskEnvelope]
+    result: type[TaskResultEnvelope]
+
+
+TASK_MODELS: Mapping[TaskType, TaskModels] = {
+    TaskType.CUSTOM_SQL_TEST_V1: TaskModels(
+        task=CustomSqlTestTask,
+        result=CustomSqlTestResult,
+    ),
+    TaskType.RUN_SQL_V1: TaskModels(
+        task=RunSqlTask,
+        result=RunSqlResult,
+    ),
+}

elementary_runtime/protocol/tasks/results.py

@@ -0,0 +1,30 @@
+from typing import Annotated
+
+from pydantic import AfterValidator, Field
+
+from elementary_runtime.common.validators import require_unique_attr
+from elementary_runtime.protocol.base import IdStr, ProtocolModel
+from elementary_runtime.protocol.enums import SubmittedResultStatus
+from elementary_runtime.protocol.limits import MAX_RESULTS_PER_BATCH
+from elementary_runtime.protocol.tasks.registry import TaskResult
+
+
+class SubmitResultsRequest(ProtocolModel):
+    results: Annotated[
+        list[TaskResult],
+        Field(max_length=MAX_RESULTS_PER_BATCH),
+        AfterValidator(require_unique_attr("task_id")),
+    ]
+
+
+class SubmittedResult(ProtocolModel):
+    task_id: IdStr
+    status: SubmittedResultStatus
+
+
+class SubmitResultsResponse(ProtocolModel):
+    results: Annotated[
+        list[SubmittedResult],
+        Field(max_length=MAX_RESULTS_PER_BATCH),
+        AfterValidator(require_unique_attr("task_id")),
+    ]

elementary_runtime/protocol/tasks/run_sql.py

@@ -0,0 +1,29 @@
+from typing import Literal
+
+from elementary_runtime.protocol.enums import TaskType
+from elementary_runtime.protocol.tasks.envelope import (
+    TaskResultEnvelope,
+    WarehouseTaskEnvelope,
+)
+from elementary_runtime.protocol.tasks.sql_query import (
+    SqlQueryResultPayload,
+    SqlQueryTaskPayload,
+)
+
+
+class RunSqlTaskPayload(SqlQueryTaskPayload):
+    pass
+
+
+class RunSqlTask(WarehouseTaskEnvelope):
+    task_type: Literal[TaskType.RUN_SQL_V1]
+    payload: RunSqlTaskPayload
+
+
+class RunSqlResultPayload(SqlQueryResultPayload):
+    pass
+
+
+class RunSqlResult(TaskResultEnvelope):
+    task_type: Literal[TaskType.RUN_SQL_V1]
+    payload: RunSqlResultPayload

elementary_runtime/protocol/tasks/sql_query.py

@@ -0,0 +1,102 @@
+from typing import Annotated
+
+from pydantic import AfterValidator, Field, JsonValue, model_validator
+from typing_extensions import Self
+
+from elementary_runtime.common.validators import require_unique_attr
+from elementary_runtime.protocol.base import (
+    IdStr,
+    NonEmptyStr,
+    ProtocolEnum,
+    ProtocolModel,
+)
+from elementary_runtime.protocol.errors import TaskErrorDetail
+from elementary_runtime.protocol.limits import (
+    MAX_QUERY_TIMEOUT_SECONDS,
+    MAX_SAMPLE_COLUMNS,
+    MAX_SAMPLE_ROW_LIMIT,
+    MAX_SAMPLE_ROWS,
+)
+
+
+class SqlExecutionStatus(ProtocolEnum):
+    COMPLETED = "completed"
+    SQL_VALIDATION_FAILED = "sql_validation_failed"
+    QUERY_ERROR = "query_error"
+    QUERY_TIMEOUT = "query_timeout"
+    PERMISSION_DENIED = "permission_denied"
+    TIMEOUT_NOT_SUPPORTED = "timeout_not_supported"
+    RESULT_LIMIT_EXCEEDED = "result_limit_exceeded"
+
+
+class SampleColumnType(ProtocolEnum):
+    STR = "str"
+    INT = "int"
+    FLOAT = "float"
+    BOOL = "bool"
+    DATE = "date"
+    TIMESTAMP = "timestamp"
+
+
+class SampleHeader(ProtocolModel):
+    id: IdStr
+    display_name: NonEmptyStr
+    type: SampleColumnType
+
+
+class ResultSample(ProtocolModel):
+    headers: Annotated[
+        list[SampleHeader],
+        Field(max_length=MAX_SAMPLE_COLUMNS),
+        AfterValidator(require_unique_attr("id")),
+    ]
+    rows: Annotated[list[dict[str, JsonValue]], Field(max_length=MAX_SAMPLE_ROWS)]
+
+    @model_validator(mode="after")
+    def rows_match_headers(self) -> Self:
+        expected = {header.id for header in self.headers}
+        for index, row in enumerate(self.rows):
+            if set(row) != expected:
+                raise ValueError(f"row {index} keys must equal the header ids")
+        return self
+
+
+class SqlQueryTaskPayload(ProtocolModel):
+    sql: Annotated[str, Field(min_length=1)]
+    query_timeout_seconds: Annotated[int, Field(ge=1, le=MAX_QUERY_TIMEOUT_SECONDS)]
+    sample_row_limit: Annotated[int, Field(ge=0, le=MAX_SAMPLE_ROW_LIMIT)]
+
+
+class SqlQueryResultPayload(ProtocolModel):
+    execution_status: SqlExecutionStatus
+    row_count: Annotated[int, Field(ge=0)] | None
+    sample: ResultSample | None
+    error: TaskErrorDetail | None
+
+    @model_validator(mode="after")
+    def status_fields_are_consistent(self) -> Self:
+        if self.execution_status is SqlExecutionStatus.COMPLETED:
+            if self.row_count is None:
+                raise ValueError(
+                    "row_count is required when execution_status is completed"
+                )
+            if self.error is not None:
+                raise ValueError(
+                    "error must be null when execution_status is completed"
+                )
+            if self.sample is not None and len(self.sample.rows) > self.row_count:
+                raise ValueError("sample cannot contain more rows than row_count")
+        else:
+            if self.error is None:
+                raise ValueError(
+                    "error is required when execution_status is not completed"
+                )
+            if self.row_count is not None:
+                raise ValueError(
+                    "row_count must be null when execution_status is not completed"
+                )
+            if self.sample is not None:
+                raise ValueError(
+                    "sample must be null when execution_status is not completed"
+                )
+        return self