use std::collections::BTreeMap;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::DataTenantId;
use crate::auth::{Attribution, PrincipalId, PrincipalKindTag};
use crate::error::WyrdProblem;
use crate::reference::CardRef;
use crate::request_id::RequestId;
pub use crate::vala::audit_detail::{
AuditDelegationStep, AuditDetail, AuditDetailValueError, BatchId, BifrostSecurityPhase,
BifrostSecurityViolationKind, ForgeIcebergRewritePhase, ForgeManifestRewritePhase,
ForgeOrphanGcPhase, ForgePromotedFile, ForgePromotedFileSetDigest, ForgeScribePromotionPhase,
ForgeSnapshotExpirePhase, OracleTableProtectionPhase, QueryAuditDigest, QueryExecutionMode,
ScopeHash, StoragePath, audit_detail_canonical_json,
};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(transparent)]
pub struct BifrostTableName(String);
impl BifrostTableName {
#[must_use]
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for BifrostTableName {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
fn default_true() -> bool {
true
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub enum TableStatus {
Active,
Deprecated,
Quarantined,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct BifrostTableEntry {
pub namespace: String,
pub name: String,
pub table_uid: String,
pub status: TableStatus,
pub fingerprint: String,
pub registered_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct BifrostTableDescription {
pub entry: BifrostTableEntry,
pub user_fields: Vec<FieldSpec>,
pub correlation_fields: Vec<FieldSpec>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schemars(schema_with = "crate::schema::skipped_option::<String>")]
pub canonical_physical_fingerprint: Option<String>,
pub physical_layout: PhysicalLayoutWire,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schemars(schema_with = "crate::schema::skipped_option::<u64>")]
pub compaction_target_file_size_bytes: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schemars(schema_with = "crate::schema::skipped_option::<CompactionTypeWire>")]
pub compaction_type: Option<CompactionTypeWire>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "kebab-case")]
pub enum CompactionTypeWire {
Auto,
Full,
SmallFiles,
FilesWithDelete,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub enum TimeUnit {
Second,
Millisecond,
Microsecond,
Nanosecond,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub enum DataTypeSpec {
Bool,
Int8,
Int16,
Int32,
Int64,
UInt8,
UInt16,
UInt32,
UInt64,
Float32,
Float64,
Utf8,
LargeUtf8,
Binary,
LargeBinary,
FixedSizeBinary {
len: i32,
},
Date32,
Date64,
Timestamp {
unit: TimeUnit,
#[serde(default)]
tz: Option<String>,
},
Time32 {
unit: TimeUnit,
},
Time64 {
unit: TimeUnit,
},
Decimal128 {
precision: u8,
scale: i8,
},
#[cfg_attr(feature = "server", schema(no_recursion))]
List(Box<FieldSpec>),
#[cfg_attr(feature = "server", schema(no_recursion))]
Struct(Vec<FieldSpec>),
#[cfg_attr(feature = "server", schema(no_recursion))]
Map {
key: Box<FieldSpec>,
value: Box<FieldSpec>,
},
Variant,
}
pub const VARIANT_MAX_DEPTH: u32 = 64;
pub const VARIANT_MAX_ENCODED_BYTES: u64 = 8_388_608;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct FieldSpec {
pub name: String,
pub data_type: DataTypeSpec,
#[serde(default = "default_true")]
pub nullable: bool,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub metadata: BTreeMap<String, String>,
}
pub const PARQUET_FIELD_ID_KEY: &str = "PARQUET:field_id";
pub const UTC_TIME_ZONE: &str = "+00:00";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum TimeGranularityWire {
Hour,
Day,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum SortDirectionWire {
Asc,
Desc,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum NullOrderWire {
First,
Last,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct SortKeyWire {
pub column: String,
pub direction: SortDirectionWire,
pub null_order: NullOrderWire,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct PhysicalLayoutWire {
pub partition_granularity: TimeGranularityWire,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schemars(schema_with = "crate::schema::skipped_option::<String>")]
pub time_column: Option<String>,
#[serde(default)]
pub sort_keys: Vec<SortKeyWire>,
#[serde(default)]
pub bloom_columns: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct RegisterTableRequest {
pub namespace: String,
pub name: String,
pub fields: Vec<FieldSpec>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub physical_layout: Option<PhysicalLayoutWire>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compaction_target_file_size_bytes: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compaction_type: Option<CompactionTypeWire>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub enum RegisterOutcome {
Created,
AlreadyExists,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct RegisterTableResponse {
pub outcome: RegisterOutcome,
pub table_uid: String,
pub fingerprint: String,
}
pub const MAX_QUERY_TERMINAL_WARNINGS: usize = 16;
pub const MAX_QUERY_SOURCE_COMPLETIONS: usize = 3;
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum QueryContractError {
#[error("{field} must not be empty")]
Empty {
field: &'static str,
},
#[error("{field} exceeds its maximum of {maximum}")]
TooMany {
field: &'static str,
maximum: usize,
},
#[error("invalid query terminal: {reason}")]
InvalidTerminal {
reason: &'static str,
},
#[error("params[{index}] must be a finite number")]
InvalidParam {
index: usize,
},
}
impl From<QueryContractError> for crate::vala::error::BifrostError {
fn from(error: QueryContractError) -> Self {
Self::QueryInvalidSql {
detail: error.to_string(),
}
}
}
impl From<QueryContractError> for crate::error::WyrdError {
fn from(error: QueryContractError) -> Self {
Self::Vala {
error: error.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(untagged)]
pub enum QueryParam {
Null,
Bool(bool),
Int(i64),
Float(f64),
String(String),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct BifrostQueryRequest {
pub sql: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub params: Vec<QueryParam>,
#[schemars(range(min = 1, max = 4_294_967_295_u32))]
#[cfg_attr(feature = "server", schema(minimum = 1, maximum = 4_294_967_295_u32))]
pub deadline_ms: Option<i64>,
}
impl BifrostQueryRequest {
pub fn validate(&self) -> Result<(), QueryContractError> {
if self.sql.trim().is_empty() {
return Err(QueryContractError::Empty { field: "sql" });
}
if let Some(index) = self
.params
.iter()
.position(|param| matches!(param, QueryParam::Float(value) if !value.is_finite()))
{
return Err(QueryContractError::InvalidParam { index });
}
if self
.deadline_ms
.is_some_and(|deadline| !(1..=i64::from(u32::MAX)).contains(&deadline))
{
return Err(QueryContractError::InvalidTerminal {
reason: "deadline_ms must be between 1 and 4294967295",
});
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum QueryClass {
Interactive,
Analytical,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum RunningQueryLifecycleState {
Admitted,
Running,
Cancelling,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct RunningQueryProgress {
pub completed_participants: u32,
pub total_participants: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct RunningQuerySummary {
pub request_id: RequestId,
pub query_class: QueryClass,
pub started_at: DateTime<Utc>,
pub deadline: DateTime<Utc>,
pub state: RunningQueryLifecycleState,
pub progress: RunningQueryProgress,
pub cancellation_requested: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ListRunningQueriesResponse {
pub queries: Vec<RunningQuerySummary>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct GetRunningQueryRequest {
pub request_id: RequestId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct CancelRunningQueryRequest {
pub request_id: RequestId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct CancelRunningQueryResponse {
pub request_id: RequestId,
pub cancellation_started: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct OracleRoleFence {
pub node_id: NodeId,
pub role: ClusterRole,
pub fencing_token: FencingToken,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct OracleLifecycleLookupRequest {
pub tenant_id: DataTenantId,
pub request_id: RequestId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ListOracleLifecyclesRequest {
pub tenant_id: DataTenantId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ListOracleLifecyclesResponse {
pub queries: Vec<RunningQuerySummary>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct GetOracleLifecycleResponse {
pub query: RunningQuerySummary,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct CancelOracleLifecycleRequest {
pub tenant_id: DataTenantId,
pub request_id: RequestId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct CancelOracleLifecycleResponse {
pub request_id: RequestId,
pub cancellation_started: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(tag = "kind", content = "data", rename_all = "snake_case")]
pub enum QueryStreamFrame {
Schema(QuerySchemaFrame),
Batch(QueryBatchFrame),
Terminal(QueryTerminalFrame),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct QuerySchemaFrame {
pub schema_fingerprint: String,
pub arrow_ipc_schema: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct QueryBatchFrame {
pub arrow_ipc_batch: Vec<u8>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum QueryTerminalOutcome {
Success,
Degraded,
Failed,
}
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, schemars::JsonSchema,
)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum QuerySource {
Iceberg,
HotSealed,
LiveTail,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum QueryWarning {
LiveTailUnavailable,
StaleCutReplanned,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum SourceCompletionOutcome {
Complete,
Unavailable,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct SourceCompletion {
pub source: QuerySource,
pub outcome: SourceCompletionOutcome,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct QueryTerminalFrame {
pub outcome: QueryTerminalOutcome,
pub query_class: QueryClass,
pub row_count: u64,
pub warnings: Vec<QueryWarning>,
pub source_completion: Vec<SourceCompletion>,
#[serde(default)]
pub error: Option<Box<WyrdProblem>>,
pub arrow_ipc_eos: Vec<u8>,
}
impl QueryTerminalFrame {
pub fn validate(&self) -> Result<(), QueryContractError> {
if self.warnings.len() > MAX_QUERY_TERMINAL_WARNINGS {
return Err(QueryContractError::TooMany {
field: "warnings",
maximum: MAX_QUERY_TERMINAL_WARNINGS,
});
}
if self.source_completion.len() != MAX_QUERY_SOURCE_COMPLETIONS {
return Err(QueryContractError::InvalidTerminal {
reason: "source completion must name every source tier",
});
}
let mut seen = std::collections::BTreeSet::new();
if self
.source_completion
.iter()
.any(|entry| !seen.insert(entry.source))
{
return Err(QueryContractError::InvalidTerminal {
reason: "source completion contains duplicates",
});
}
for source in [QuerySource::Iceberg, QuerySource::HotSealed] {
if !self.source_completion.iter().any(|entry| {
entry.source == source && entry.outcome == SourceCompletionOutcome::Complete
}) {
return Err(QueryContractError::InvalidTerminal {
reason: "sealed sources must be complete",
});
}
}
let degraded_live = self.source_completion.iter().any(|entry| {
entry.source == QuerySource::LiveTail
&& entry.outcome == SourceCompletionOutcome::Unavailable
});
let failed = self.outcome == QueryTerminalOutcome::Failed;
if failed != self.error.is_some() {
return Err(QueryContractError::InvalidTerminal {
reason: "error presence must match failed outcome",
});
}
if !failed && degraded_live != (self.outcome == QueryTerminalOutcome::Degraded) {
return Err(QueryContractError::InvalidTerminal {
reason: "degraded outcome must match unavailable live source",
});
}
if self.warnings.contains(&QueryWarning::LiveTailUnavailable) != degraded_live {
return Err(QueryContractError::InvalidTerminal {
reason: "live-tail warning must match unavailable source",
});
}
if failed != self.arrow_ipc_eos.is_empty() {
return Err(QueryContractError::InvalidTerminal {
reason: if failed {
"failed terminal must not carry an Arrow IPC end-of-stream"
} else {
"successful terminal must carry its Arrow IPC end-of-stream"
},
});
}
Ok(())
}
pub fn validate_emitted_rows(&self, emitted_rows: u64) -> Result<(), QueryContractError> {
if self.row_count != emitted_rows {
return Err(QueryContractError::InvalidTerminal {
reason: "terminal row count does not match emitted rows",
});
}
Ok(())
}
}
#[cfg(test)]
mod query_terminal_tests {
use super::*;
use QueryTerminalOutcome::{Degraded, Failed, Success};
use SourceCompletionOutcome::{Complete, Unavailable};
const EOS: [u8; 8] = [0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0];
fn sources(live: SourceCompletionOutcome) -> Vec<SourceCompletion> {
vec![
SourceCompletion {
source: QuerySource::Iceberg,
outcome: SourceCompletionOutcome::Complete,
},
SourceCompletion {
source: QuerySource::HotSealed,
outcome: SourceCompletionOutcome::Complete,
},
SourceCompletion {
source: QuerySource::LiveTail,
outcome: live,
},
]
}
fn execution_failed() -> Box<WyrdProblem> {
Box::new(
crate::error::WyrdError::from(crate::vala::error::BifrostError::QueryExecutionFailed)
.problem(),
)
}
fn terminal(
outcome: QueryTerminalOutcome,
live: SourceCompletionOutcome,
warnings: Vec<QueryWarning>,
) -> QueryTerminalFrame {
let failed = outcome == QueryTerminalOutcome::Failed;
QueryTerminalFrame {
outcome,
query_class: QueryClass::Interactive,
row_count: 1,
warnings,
source_completion: sources(live),
error: failed.then(execution_failed),
arrow_ipc_eos: if failed { Vec::new() } else { EOS.to_vec() },
}
}
#[test]
fn closed_terminal_matrix_validates() {
let warn = || vec![QueryWarning::LiveTailUnavailable];
let success = terminal(Success, Complete, vec![]);
success.validate().expect("success terminal validates");
success
.validate_emitted_rows(1)
.expect("success row count validates");
let encoded = serde_json::to_value(&success).expect("terminal serializes");
assert!(encoded.get("freshness").is_none(), "no freshness field");
assert_eq!(
encoded["source_completion"].as_array().map(Vec::len),
Some(3)
);
terminal(Degraded, Unavailable, warn())
.validate()
.expect("degraded terminal validates");
terminal(Failed, Unavailable, warn())
.validate()
.expect("failure after known live loss validates");
terminal(Failed, Complete, vec![])
.validate()
.expect("failure with complete live validates");
for invalid in [
terminal(Success, Unavailable, warn()),
terminal(Success, Complete, warn()),
terminal(Degraded, Complete, vec![]),
terminal(Degraded, Unavailable, vec![]),
terminal(Failed, Unavailable, vec![]),
] {
assert!(invalid.validate().is_err(), "{invalid:?} must be rejected");
}
let mut two_tiers = terminal(Success, Complete, vec![]);
two_tiers.source_completion.pop();
assert!(two_tiers.validate().is_err(), "every tier is required");
let mut unavailable_published = terminal(Failed, Complete, vec![]);
unavailable_published.source_completion[0].outcome = Unavailable;
assert!(
unavailable_published.validate().is_err(),
"published tiers are never degraded"
);
}
#[test]
fn invalid_failed_terminal_is_rejected() {
let mut failed_without_error = terminal(
QueryTerminalOutcome::Failed,
SourceCompletionOutcome::Complete,
vec![],
);
failed_without_error.error = None;
assert!(failed_without_error.validate().is_err());
let failed = terminal(
QueryTerminalOutcome::Failed,
SourceCompletionOutcome::Complete,
vec![],
);
assert!(failed.validate_emitted_rows(2).is_err());
}
#[test]
fn query_request_has_no_source_selectors_and_validation_is_closed() {
let request: BifrostQueryRequest = serde_json::from_value(serde_json::json!({
"sql": "SELECT 1",
"deadline_ms": null
}))
.expect("request deserializes");
request.validate().expect("minimal request validates");
assert_eq!(
serde_json::to_value(&request).expect("request serializes"),
serde_json::json!({"sql": "SELECT 1", "deadline_ms": null})
);
for (field, value) in [
("visibility", "published_only"),
("freshness", "allow_degraded"),
("query_class", "analytical"),
] {
let mut body = serde_json::json!({"sql": "SELECT 1", "deadline_ms": null});
body[field] = serde_json::Value::from(value);
assert!(
serde_json::from_value::<BifrostQueryRequest>(body).is_err(),
"{field} must be rejected"
);
}
let invalid = BifrostQueryRequest {
params: Vec::new(),
sql: " ".into(),
deadline_ms: Some(0),
};
assert!(invalid.validate().is_err());
}
#[test]
fn query_request_deadline_range_is_closed_and_published() {
let request = |deadline_ms| BifrostQueryRequest {
params: Vec::new(),
sql: "SELECT 1".into(),
deadline_ms: Some(deadline_ms),
};
let max = i64::from(u32::MAX);
for valid in [1, max] {
request(valid)
.validate()
.expect("in-range deadline validates");
}
for invalid in [-1, 0, max + 1] {
assert!(
request(invalid).validate().is_err(),
"{invalid} must be rejected"
);
}
let decoded: BifrostQueryRequest = serde_json::from_value(serde_json::json!({
"sql": "SELECT 1",
"deadline_ms": max + 1
}))
.expect("an above-range deadline deserializes for validation");
assert!(decoded.validate().is_err());
let schema = serde_json::to_value(schemars::schema_for!(BifrostQueryRequest).schema)
.expect("request schema serializes");
let deadline = &schema["properties"]["deadline_ms"];
assert_eq!(deadline["minimum"].as_f64(), Some(1.0));
assert_eq!(deadline["maximum"].as_f64(), Some(f64::from(u32::MAX)));
assert_eq!(
schema["additionalProperties"],
serde_json::Value::Bool(false)
);
}
#[test]
fn query_request_refusals_project_onto_their_catalog_codes() {
let request = |sql: &str, deadline_ms| BifrostQueryRequest {
params: Vec::new(),
sql: sql.into(),
deadline_ms,
};
let refusal = |request: BifrostQueryRequest| {
crate::error::WyrdError::from(request.validate().expect_err("request is refused"))
.code()
};
assert_eq!(
refusal(request("SELECT 1", Some(0))),
"WYRD_VALA_400_QUERY_INVALID_SQL"
);
assert_eq!(
refusal(request(" ", None)),
"WYRD_VALA_400_QUERY_INVALID_SQL"
);
}
}
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Serialize,
Deserialize,
schemars::JsonSchema,
)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(transparent)]
pub struct NodeId(
uuid::Uuid,
);
impl NodeId {
#[must_use]
pub const fn new(value: uuid::Uuid) -> Self {
Self(value)
}
#[must_use]
pub const fn as_uuid(self) -> uuid::Uuid {
self.0
}
}
impl From<NodeId> for uuid::Uuid {
fn from(value: NodeId) -> Self {
value.0
}
}
impl From<uuid::Uuid> for NodeId {
fn from(value: uuid::Uuid) -> Self {
Self(value)
}
}
pub type FencingToken = u64;
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Serialize,
Deserialize,
schemars::JsonSchema,
)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(transparent)]
pub struct QueryId(
uuid::Uuid,
);
impl QueryId {
#[must_use]
pub const fn new(value: uuid::Uuid) -> Self {
Self(value)
}
#[must_use]
pub const fn as_uuid(self) -> uuid::Uuid {
self.0
}
}
impl From<QueryId> for uuid::Uuid {
fn from(value: QueryId) -> Self {
value.0
}
}
impl From<uuid::Uuid> for QueryId {
fn from(value: uuid::Uuid) -> Self {
Self(value)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum ClusterRole {
Scribe,
Oracle,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ClusterNodeKey {
pub node_id: NodeId,
pub role: ClusterRole,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ScribeCapabilitiesV1 {
pub tail_protocol_version: u16,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct OracleCapabilitiesV1 {
pub storage_protocol_version: u16,
pub cpu_cores: f64,
pub memory_budget_bytes: u64,
pub cpu_cores_per_slot: f64,
pub memory_bytes_per_slot: u64,
pub raw_slots: u32,
pub usable_slots: u32,
pub supported_classes: Vec<QueryClass>,
pub max_workers_per_query: u32,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ClusterCapabilities {
ScribeV1(ScribeCapabilitiesV1),
OracleV1(OracleCapabilitiesV1),
}
impl ClusterCapabilities {
pub fn validate_for_role(&self, role: ClusterRole) -> Result<(), QueryContractError> {
match (role, self) {
(ClusterRole::Scribe, Self::ScribeV1(value)) if value.tail_protocol_version == 1 => {
Ok(())
}
(ClusterRole::Oracle, Self::OracleV1(value))
if value.storage_protocol_version == 1
&& value.cpu_cores.is_finite()
&& value.cpu_cores > 0.0
&& value.cpu_cores_per_slot.is_finite()
&& value.cpu_cores_per_slot > 0.0
&& value.memory_budget_bytes > 0
&& value.memory_bytes_per_slot > 0
&& value.raw_slots > 0
&& value.usable_slots > 0
&& value.usable_slots <= value.raw_slots
&& value.max_workers_per_query <= 63
&& !value.supported_classes.is_empty() =>
{
Ok(())
}
_ => Err(QueryContractError::InvalidTerminal {
reason: "invalid role capability document",
}),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ClusterRoleLease {
pub key: ClusterNodeKey,
pub address: String,
pub fencing_token: FencingToken,
pub capability_version: u16,
pub capabilities: ClusterCapabilities,
pub ready: bool,
pub started_at: DateTime<Utc>,
pub heartbeat_at: DateTime<Utc>,
}
pub type WriterEpoch = u64;
pub type WalLsn = u64;
macro_rules! private_uuid_id {
($name:ident, $doc:literal) => {
#[doc = $doc]
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
PartialOrd,
Ord,
Hash,
Serialize,
Deserialize,
schemars::JsonSchema,
)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(transparent)]
pub struct $name(
/// Opaque UUID carried by the private control-plane contract.
uuid::Uuid,
);
impl $name {
#[must_use]
pub const fn new(value: uuid::Uuid) -> Self {
Self(value)
}
#[must_use]
pub const fn as_uuid(self) -> uuid::Uuid {
self.0
}
}
};
}
private_uuid_id!(
ReservationId,
"Opaque identity for one pending Oracle worker reservation."
);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct TimePartitionWire {
granularity: TimeGranularityWire,
start_utc: DateTime<Utc>,
}
impl TimePartitionWire {
pub fn new(
granularity: TimeGranularityWire,
start_utc: DateTime<Utc>,
) -> Result<Self, QueryContractError> {
if !partition_start_is_canonical(granularity, start_utc) {
return Err(QueryContractError::InvalidTerminal {
reason: "time partition start must be the exact UTC boundary of its granularity",
});
}
Ok(Self {
granularity,
start_utc,
})
}
#[must_use]
pub const fn granularity(&self) -> TimeGranularityWire {
self.granularity
}
#[must_use]
pub const fn start_utc(&self) -> DateTime<Utc> {
self.start_utc
}
#[must_use]
pub fn start_unix_micros(&self) -> i64 {
self.start_utc.timestamp_micros()
}
#[must_use]
pub const fn granularity_tag(&self) -> u8 {
match self.granularity {
TimeGranularityWire::Hour => 1,
TimeGranularityWire::Day => 2,
}
}
#[must_use]
pub fn as_path_components(&self) -> String {
format!(
"partition_granularity={}/partition_start={}",
self.granularity_str(),
self.start_utc.format("%Y-%m-%dT%HZ")
)
}
#[must_use]
pub const fn granularity_str(&self) -> &'static str {
match self.granularity {
TimeGranularityWire::Hour => "hour",
TimeGranularityWire::Day => "day",
}
}
}
impl PartialOrd for TimePartitionWire {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for TimePartitionWire {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.granularity_tag()
.cmp(&other.granularity_tag())
.then(self.start_utc.cmp(&other.start_utc))
}
}
impl std::fmt::Display for TimePartitionWire {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}:{}",
self.granularity_str(),
self.start_utc.to_rfc3339()
)
}
}
impl<'de> Deserialize<'de> for TimePartitionWire {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(Deserialize)]
struct Raw {
granularity: TimeGranularityWire,
start_utc: DateTime<Utc>,
}
let raw = Raw::deserialize(deserializer)?;
Self::new(raw.granularity, raw.start_utc).map_err(serde::de::Error::custom)
}
}
#[must_use]
pub fn partition_start_is_canonical(
granularity: TimeGranularityWire,
start_utc: DateTime<Utc>,
) -> bool {
use chrono::Timelike as _;
let sub_hour_is_zero =
start_utc.minute() == 0 && start_utc.second() == 0 && start_utc.nanosecond() == 0;
match granularity {
TimeGranularityWire::Hour => sub_hour_is_zero,
TimeGranularityWire::Day => sub_hour_is_zero && start_utc.hour() == 0,
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(transparent)]
pub struct SchemaFingerprint(
String,
);
impl SchemaFingerprint {
pub fn new(value: impl Into<String>) -> Result<Self, QueryContractError> {
let value = value.into();
if value.is_empty() || value.len() > 1_024 {
return Err(QueryContractError::InvalidTerminal {
reason: "invalid schema fingerprint",
});
}
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl<'de> Deserialize<'de> for SchemaFingerprint {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
Self::new(String::deserialize(deserializer)?).map_err(serde::de::Error::custom)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct TenantTableBinding {
pub tenant_id: crate::DataTenantId,
pub namespace: String,
pub table: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct TailStreamIdentity {
pub node_id: NodeId,
pub writer_epoch: WriterEpoch,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ReserveNodeSlotsRequest {
pub query_id: QueryId,
pub leader_node_id: NodeId,
pub leader_fencing_token: FencingToken,
pub expires_at: DateTime<Utc>,
pub graph: AnalyticalGraphRef,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct AnalyticalGraphRef {
pub public_query_id: uuid::Uuid,
pub datafusion_query_id: uuid::Uuid,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct PendingNodeReservation {
pub reservation_id: ReservationId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ReservationRejected {
pub retry_after_ms: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub enum ReserveNodeSlotsResponse {
Pending(PendingNodeReservation),
Rejected(ReservationRejected),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct PeerContext {
pub claims_bytes: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ScribeProviderCut {
pub writer_epoch: u64,
pub start_partition: TimePartitionWire,
pub end_partition: TimePartitionWire,
}
impl ScribeProviderCut {
#[must_use]
pub fn is_valid(&self) -> bool {
self.writer_epoch > 0
&& self.start_partition.granularity() == self.end_partition.granularity()
&& self.start_partition <= self.end_partition
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct FollowerScanAssignment {
pub scan_id: String,
pub binding: TenantTableBinding,
pub scribe_provider_cut: Option<ScribeProviderCut>,
pub schema_fingerprint: String,
pub required_columns: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct ExecuteFragmentRequest {
pub context: PeerContext,
pub physical_plan_bytes: Vec<u8>,
pub reservation_id: ReservationId,
pub leader_fence: OracleRoleFence,
pub target_fence: OracleRoleFence,
pub assignments: Vec<FollowerScanAssignment>,
pub plan_fingerprint: String,
}
#[derive(
Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema,
)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct WorkerScanStats {
pub bytes_scanned: Option<u64>,
pub files_scanned: u64,
pub partitions_scanned: u64,
pub row_groups_scanned: u64,
pub row_groups_pruned: u64,
pub row_groups_pruned_bloom: u64,
pub rows_pruned_page_index: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct WorkerFooter {
pub fragment_id: String,
pub manifest_digest: QueryAuditDigest,
pub row_count: u64,
pub encoded_bytes: u64,
pub payload_digest: QueryAuditDigest,
pub completed: bool,
pub scan_stats: WorkerScanStats,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub enum WorkerAttemptFrame {
Schema(Vec<u8>),
Batch(Vec<u8>),
Footer(WorkerFooter),
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn query_bind_values_round_trip_in_order() {
let body = serde_json::json!({
"sql": "SELECT * FROM t WHERE a = $1 AND b = $2 AND c = $3 AND d = $4 AND e = $5",
"params": [null, true, 42, 1.5, "x'; DROP TABLE t; --"],
"deadline_ms": null
});
let request: BifrostQueryRequest =
serde_json::from_value(body.clone()).expect("request deserializes");
assert_eq!(
request.params,
[
QueryParam::Null,
QueryParam::Bool(true),
QueryParam::Int(42),
QueryParam::Float(1.5),
QueryParam::String("x'; DROP TABLE t; --".into()),
]
);
request.validate().expect("finite scalars validate");
assert_eq!(
serde_json::to_value(&request).expect("request serializes"),
body
);
let invalid = BifrostQueryRequest {
params: vec![QueryParam::Int(1), QueryParam::Float(f64::NAN)],
..request
};
assert_eq!(
invalid.validate(),
Err(QueryContractError::InvalidParam { index: 1 })
);
}
#[test]
fn query_terminal_eos_contract() {
const EOS: [u8; 8] = [0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0];
let sources = vec![
SourceCompletion {
source: QuerySource::Iceberg,
outcome: SourceCompletionOutcome::Complete,
},
SourceCompletion {
source: QuerySource::HotSealed,
outcome: SourceCompletionOutcome::Complete,
},
SourceCompletion {
source: QuerySource::LiveTail,
outcome: SourceCompletionOutcome::Complete,
},
];
let base = QueryTerminalFrame {
outcome: QueryTerminalOutcome::Success,
query_class: QueryClass::Interactive,
row_count: 0,
warnings: vec![],
source_completion: sources,
error: None,
arrow_ipc_eos: EOS.to_vec(),
};
base.validate()
.expect("empty success carries its end-of-stream");
base.validate_emitted_rows(0)
.expect("empty success emitted no rows");
assert!(
QueryTerminalFrame {
arrow_ipc_eos: Vec::new(),
..base.clone()
}
.validate()
.is_err(),
"success without an end-of-stream must be rejected"
);
let mut degraded_sources = base.source_completion.clone();
degraded_sources[2].outcome = SourceCompletionOutcome::Unavailable;
let degraded = QueryTerminalFrame {
outcome: QueryTerminalOutcome::Degraded,
warnings: vec![QueryWarning::LiveTailUnavailable],
source_completion: degraded_sources,
..base.clone()
};
degraded
.validate()
.expect("degraded still finishes its Arrow stream");
assert!(
QueryTerminalFrame {
arrow_ipc_eos: Vec::new(),
..degraded
}
.validate()
.is_err(),
"degraded without an end-of-stream must be rejected"
);
let failed = QueryTerminalFrame {
outcome: QueryTerminalOutcome::Failed,
error: Some(Box::new(
crate::error::WyrdError::from(
crate::vala::error::BifrostError::QueryExecutionFailed,
)
.problem(),
)),
arrow_ipc_eos: Vec::new(),
..base.clone()
};
failed.validate().expect("a failed terminal closes nothing");
assert!(
QueryTerminalFrame {
arrow_ipc_eos: EOS.to_vec(),
..failed
}
.validate()
.is_err(),
"a failed terminal must not claim an end-of-stream"
);
let encoded = serde_json::to_value(&base).expect("terminal serializes");
assert_eq!(
serde_json::from_value::<QueryTerminalFrame>(encoded).expect("terminal deserializes"),
base
);
}
#[test]
fn query_terminal_projects_path_without_request_selector() {
let request = serde_json::to_value(schemars::schema_for!(BifrostQueryRequest).schema)
.expect("request schema serializes");
let mut request_properties = request["properties"]
.as_object()
.expect("request has properties")
.keys()
.map(String::as_str)
.collect::<Vec<_>>();
request_properties.sort_unstable();
assert_eq!(
request_properties,
["deadline_ms", "params", "sql"],
"the request must not accept a source, freshness, path, class, worker, or plan selector"
);
let path = serde_json::to_value(schemars::schema_for!(QueryClass).schema)
.expect("path schema serializes");
let path_variants = path["oneOf"]
.as_array()
.expect("the selected path is a closed enum")
.iter()
.map(|variant| {
variant["enum"][0]
.as_str()
.expect("each variant is one string constant")
.to_owned()
})
.collect::<Vec<_>>();
assert_eq!(
path_variants,
["interactive", "analytical"],
"the selected path is closed over Interactive and Analytical"
);
let terminal = QueryTerminalFrame {
outcome: QueryTerminalOutcome::Success,
query_class: QueryClass::Analytical,
row_count: 0,
warnings: vec![],
source_completion: vec![
SourceCompletion {
source: QuerySource::Iceberg,
outcome: SourceCompletionOutcome::Complete,
},
SourceCompletion {
source: QuerySource::HotSealed,
outcome: SourceCompletionOutcome::Complete,
},
SourceCompletion {
source: QuerySource::LiveTail,
outcome: SourceCompletionOutcome::Complete,
},
],
error: None,
arrow_ipc_eos: vec![0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0],
};
terminal
.validate()
.expect("an Analytical terminal is valid");
let encoded = serde_json::to_value(&terminal).expect("terminal serializes");
assert_eq!(
encoded["query_class"],
serde_json::Value::from("analytical"),
"the terminal names the server-selected path on the wire"
);
assert_eq!(
serde_json::from_value::<QueryTerminalFrame>(encoded).expect("terminal deserializes"),
terminal
);
}
#[test]
fn running_query_contract_round_trips() {
let request_id = RequestId::now_v7();
let started_at = DateTime::from_timestamp_millis(1_725_000_000_123)
.expect("fixed running-query timestamp is valid");
let summary = RunningQuerySummary {
request_id: request_id.clone(),
query_class: QueryClass::Interactive,
started_at,
deadline: started_at + chrono::Duration::seconds(30),
state: RunningQueryLifecycleState::Running,
progress: RunningQueryProgress {
completed_participants: 1,
total_participants: 2,
},
cancellation_requested: false,
};
assert_json_round_trip(&summary);
assert_json_round_trip(&ListRunningQueriesResponse {
queries: vec![summary.clone()],
});
assert_json_round_trip(&GetRunningQueryRequest {
request_id: request_id.clone(),
});
assert_json_round_trip(&CancelRunningQueryRequest {
request_id: request_id.clone(),
});
assert_json_round_trip(&CancelRunningQueryResponse {
request_id: request_id.clone(),
cancellation_started: true,
});
let tenant_id = DataTenantId::new_v7();
assert_json_round_trip(&ListOracleLifecyclesRequest { tenant_id });
assert_json_round_trip(&ListOracleLifecyclesResponse {
queries: vec![summary.clone()],
});
assert_json_round_trip(&OracleLifecycleLookupRequest {
tenant_id,
request_id: request_id.clone(),
});
assert_json_round_trip(&GetOracleLifecycleResponse {
query: summary.clone(),
});
assert_json_round_trip(&CancelOracleLifecycleRequest {
tenant_id,
request_id: request_id.clone(),
});
assert_json_round_trip(&CancelOracleLifecycleResponse {
request_id,
cancellation_started: true,
});
let _ = schemars::schema_for!(RunningQuerySummary);
let _ = schemars::schema_for!(ListRunningQueriesResponse);
let _ = schemars::schema_for!(GetRunningQueryRequest);
let _ = schemars::schema_for!(CancelRunningQueryRequest);
let _ = schemars::schema_for!(CancelRunningQueryResponse);
let _ = schemars::schema_for!(ListOracleLifecyclesRequest);
let _ = schemars::schema_for!(ListOracleLifecyclesResponse);
let _ = schemars::schema_for!(OracleLifecycleLookupRequest);
let _ = schemars::schema_for!(GetOracleLifecycleResponse);
let _ = schemars::schema_for!(CancelOracleLifecycleRequest);
let _ = schemars::schema_for!(CancelOracleLifecycleResponse);
}
fn assert_json_round_trip<T>(expected: &T)
where
T: serde::Serialize + serde::de::DeserializeOwned + std::fmt::Debug + PartialEq,
{
let value = serde_json::to_value(expected).expect("lifecycle contract serializes");
let actual: T = serde_json::from_value(value).expect("lifecycle contract deserializes");
assert_eq!(&actual, expected);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum AuthMethod {
Jwt,
Internal,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum AuditOutcome {
Allowed,
Denied,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
pub struct AuditEvent {
pub request_id: RequestId,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub trace_id: Option<String>,
pub operation: String,
pub resource: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub card_ref: Option<CardRef>,
pub sub_principal_id: PrincipalId,
pub sub_principal_kind: PrincipalKindTag,
pub act_principal_id: PrincipalId,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub credential_id: Option<Uuid>,
pub permission: String,
pub outcome: AuditOutcome,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub detail: Option<AuditDetail>,
}
impl AuditEvent {
#[must_use]
#[expect(
clippy::too_many_arguments,
reason = "the constructor mirrors the complete audit wire shape"
)]
pub fn new(
request_id: RequestId,
trace_id: Option<String>,
operation: String,
resource: String,
card_ref: Option<CardRef>,
attribution: Attribution,
sub_principal_kind: PrincipalKindTag,
permission: String,
outcome: AuditOutcome,
) -> Self {
Self {
request_id,
trace_id,
operation,
resource,
card_ref,
sub_principal_id: attribution.sub_principal_id,
sub_principal_kind,
act_principal_id: attribution.act_principal_id,
credential_id: None,
permission,
outcome,
detail: None,
}
}
#[must_use]
pub const fn with_credential_id(mut self, credential_id: Option<Uuid>) -> Self {
self.credential_id = credential_id;
self
}
#[must_use]
pub fn with_detail(mut self, detail: AuditDetail) -> Self {
self.detail = Some(detail);
self
}
}
#[cfg(test)]
mod bifrost_wire_tests {
use std::collections::BTreeMap;
use crate::vala::api::{
BifrostTableDescription, BifrostTableEntry, CompactionTypeWire, DataTypeSpec, FieldSpec,
NullOrderWire, PARQUET_FIELD_ID_KEY, PhysicalLayoutWire, RegisterOutcome,
RegisterTableRequest, RegisterTableResponse, SortDirectionWire, SortKeyWire, TableStatus,
TimeGranularityWire, TimeUnit,
};
use schemars::schema_for;
#[test]
fn async_query_contract_family_is_absent() {
let production = include_str!("api.rs")
.split("#[cfg(test)]")
.next()
.expect("Vala API has a production section");
for removed in [
"struct JobUid",
"enum AsyncJobState",
"enum ExecutorAvailability",
"struct AsyncQueryRequest",
"struct AsyncQueryResponse",
"struct AsyncQueryStatus",
] {
assert!(
!production.contains(removed),
"removed asynchronous query contract returned: {removed}"
);
}
}
fn bifrost_wire_round_trip<T>(value: &T) -> T
where
T: serde::Serialize + serde::de::DeserializeOwned + PartialEq + std::fmt::Debug,
{
let json = serde_json::to_string(value).expect("serialize");
let back: T = serde_json::from_str(&json).expect("deserialize");
assert_eq!(*value, back, "round-trip mismatch");
back
}
#[test]
fn bifrost_wire_field_spec_round_trips_and_defaults_nullable_true() {
let spec = FieldSpec {
name: "value".to_string(),
data_type: DataTypeSpec::Int64,
nullable: false,
metadata: BTreeMap::default(),
};
bifrost_wire_round_trip(&spec);
let json = serde_json::to_value(&spec).expect("serialize");
assert!(
json.get("metadata").is_none(),
"empty metadata must be skipped"
);
let minimal: FieldSpec = serde_json::from_str(r#"{"name":"x","data_type":"Utf8"}"#)
.expect("deserialize minimal");
assert!(minimal.nullable, "nullable must default to true");
assert!(minimal.metadata.is_empty());
}
#[test]
fn bifrost_wire_field_spec_carries_correlation_metadata() {
let mut spec = FieldSpec {
name: "card_uid".to_string(),
data_type: DataTypeSpec::Utf8,
nullable: true,
metadata: BTreeMap::default(),
};
spec.metadata
.insert("wyrd:column_class".to_string(), "correlation".to_string());
let back = bifrost_wire_round_trip(&spec);
assert_eq!(
back.metadata.get("wyrd:column_class").map(String::as_str),
Some("correlation")
);
}
#[test]
fn bifrost_wire_register_request_omits_physical_layout() {
let req: RegisterTableRequest =
serde_json::from_str(r#"{"namespace":"vala.bifrost","name":"events","fields":[]}"#)
.expect("deserialize");
assert!(req.physical_layout.is_none());
let json = serde_json::to_string(&req).expect("serialize");
assert!(!json.contains("physical_layout"), "{json}");
}
#[test]
fn bifrost_wire_register_request_keeps_explicit_empty_layout_lists() {
let req: RegisterTableRequest = serde_json::from_str(
r#"{"namespace":"vala.bifrost","name":"events","fields":[],
"physical_layout":{"partition_granularity":"hour",
"sort_keys":[],"bloom_columns":[]}}"#,
)
.expect("deserialize");
let layout = req.physical_layout.as_ref().expect("explicit layout");
assert_eq!(layout.partition_granularity, TimeGranularityWire::Hour);
assert!(layout.sort_keys.is_empty());
assert!(layout.bloom_columns.is_empty());
bifrost_wire_round_trip(&req);
}
#[test]
fn bifrost_wire_register_request_carries_hyphenated_compaction_type() {
let req: RegisterTableRequest = serde_json::from_str(
r#"{"namespace":"vala.datasets","name":"events","fields":[],
"compaction_type":"small-files"}"#,
)
.expect("deserialize");
assert_eq!(req.compaction_type, Some(CompactionTypeWire::SmallFiles));
bifrost_wire_round_trip(&req);
let omitted: RegisterTableRequest =
serde_json::from_str(r#"{"namespace":"vala.datasets","name":"events","fields":[]}"#)
.expect("deserialize");
assert_eq!(omitted.compaction_type, None);
assert!(
serde_json::from_str::<RegisterTableRequest>(
r#"{"namespace":"vala.datasets","name":"events","fields":[],
"compaction_type":"small_files"}"#,
)
.is_err()
);
}
#[test]
fn bifrost_wire_compaction_type_values_and_schemas_are_hyphenated() {
for (kind, spelling) in [
(CompactionTypeWire::Auto, "auto"),
(CompactionTypeWire::Full, "full"),
(CompactionTypeWire::SmallFiles, "small-files"),
(CompactionTypeWire::FilesWithDelete, "files-with-delete"),
] {
assert_eq!(
serde_json::to_value(kind).expect("serialize"),
serde_json::Value::String(spelling.to_owned())
);
assert_eq!(
serde_json::from_value::<CompactionTypeWire>(spelling.into()).expect("parse"),
kind
);
}
for refused in ["small_files", "files_with_delete"] {
assert!(serde_json::from_value::<CompactionTypeWire>(refused.into()).is_err());
}
for schema in [
serde_json::to_string(&schema_for!(RegisterTableRequest)).expect("schema"),
serde_json::to_string(&schema_for!(BifrostTableDescription)).expect("schema"),
] {
assert!(schema.contains(r#""small-files""#), "{schema}");
assert!(schema.contains(r#""files-with-delete""#), "{schema}");
assert!(!schema.contains("small_files"), "{schema}");
assert!(!schema.contains("files_with_delete"), "{schema}");
assert!(schema.contains("compacts it `small-files`"), "{schema}");
}
}
#[test]
fn bifrost_wire_physical_layout_has_no_partition_column_field() {
let layout = PhysicalLayoutWire {
partition_granularity: TimeGranularityWire::Hour,
time_column: None,
sort_keys: Vec::new(),
bloom_columns: Vec::new(),
};
let json = serde_json::to_value(&layout).expect("serialize");
let object = json.as_object().expect("layout is a JSON object");
let mut keys = object.keys().map(String::as_str).collect::<Vec<_>>();
keys.sort_unstable();
assert_eq!(
keys,
vec!["bloom_columns", "partition_granularity", "sort_keys"]
);
}
#[test]
fn bifrost_wire_nested_and_recursive_data_types_round_trip() {
let spec = FieldSpec {
name: "nested".to_string(),
data_type: DataTypeSpec::Struct(vec![
FieldSpec {
name: "tags".to_string(),
data_type: DataTypeSpec::List(Box::new(FieldSpec {
name: "item".to_string(),
data_type: DataTypeSpec::Utf8,
nullable: false,
metadata: BTreeMap::from([(
PARQUET_FIELD_ID_KEY.to_string(),
"42".to_string(),
)]),
})),
nullable: true,
metadata: BTreeMap::default(),
},
FieldSpec {
name: "ts".to_string(),
data_type: DataTypeSpec::Timestamp {
unit: TimeUnit::Microsecond,
tz: Some("UTC".to_string()),
},
nullable: false,
metadata: BTreeMap::default(),
},
]),
nullable: true,
metadata: BTreeMap::default(),
};
bifrost_wire_round_trip(&spec);
}
#[test]
fn bifrost_wire_table_entry_and_description_round_trip() {
let entry = BifrostTableEntry {
namespace: "vala.bifrost".to_string(),
name: "events".to_string(),
table_uid: "ab".repeat(16),
status: TableStatus::Active,
fingerprint: "01".repeat(32),
registered_at: chrono::Utc::now(),
updated_at: chrono::Utc::now(),
};
let desc = BifrostTableDescription {
entry: entry.clone(),
user_fields: vec![FieldSpec {
name: "value".to_string(),
data_type: DataTypeSpec::Int64,
nullable: false,
metadata: BTreeMap::default(),
}],
correlation_fields: vec![
FieldSpec {
name: "card_uid".to_string(),
data_type: DataTypeSpec::Utf8,
nullable: true,
metadata: BTreeMap::from([(
PARQUET_FIELD_ID_KEY.to_string(),
"1001".to_string(),
)]),
},
FieldSpec {
name: "run_id".to_string(),
data_type: DataTypeSpec::Utf8,
nullable: true,
metadata: BTreeMap::from([(
PARQUET_FIELD_ID_KEY.to_string(),
"1000".to_string(),
)]),
},
],
canonical_physical_fingerprint: None,
physical_layout: PhysicalLayoutWire {
partition_granularity: TimeGranularityWire::Hour,
time_column: Some("wyrd_ingested_at".to_string()),
sort_keys: vec![SortKeyWire {
column: "wyrd_ingested_at".to_string(),
direction: SortDirectionWire::Desc,
null_order: NullOrderWire::Last,
}],
bloom_columns: vec!["run_id".to_string()],
},
compaction_target_file_size_bytes: Some(1_073_741_824),
compaction_type: Some(CompactionTypeWire::SmallFiles),
};
bifrost_wire_round_trip(&entry);
bifrost_wire_round_trip(&desc);
}
#[test]
fn bifrost_wire_register_response_and_layout_round_trip() {
bifrost_wire_round_trip(&RegisterTableResponse {
outcome: RegisterOutcome::Created,
table_uid: "ab".repeat(16),
fingerprint: "01".repeat(32),
});
bifrost_wire_round_trip(&PhysicalLayoutWire {
partition_granularity: TimeGranularityWire::Day,
time_column: Some("created_at".to_string()),
sort_keys: vec![SortKeyWire {
column: "created_at".to_string(),
direction: SortDirectionWire::Desc,
null_order: NullOrderWire::Last,
}],
bloom_columns: vec!["run_id".to_string()],
});
}
#[test]
fn bifrost_wire_schema_for_wire_types_does_not_panic() {
let _ = schema_for!(BifrostTableEntry);
let _ = schema_for!(BifrostTableDescription);
let _ = schema_for!(DataTypeSpec);
let _ = schema_for!(FieldSpec);
let _ = schema_for!(RegisterTableRequest);
let _ = schema_for!(RegisterTableResponse);
}
#[test]
fn private_query_schema_types_do_not_panic() {
let _ = schema_for!(super::ReserveNodeSlotsRequest);
let _ = schema_for!(super::ReserveNodeSlotsResponse);
let _ = schema_for!(super::ExecuteFragmentRequest);
let _ = schema_for!(super::WorkerAttemptFrame);
}
}