1use crate::{
3 ProtocolVersion, RequestId,
4 asset::AssetSelector,
5 refs::{ObjectRef, SessionSelector},
6};
7use ahri_tre_types::{DisclosureBudgetOverrides, DisclosureBudgets};
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
11#[serde(deny_unknown_fields)]
12pub struct ContentRequest {
13 pub protocol_version: ProtocolVersion,
14 pub request_id: RequestId,
15 pub session: Option<SessionSelector>,
16 pub source: ContentSource,
17 pub representation: ContentRepresentation,
18 #[serde(default)]
19 pub budgets: DisclosureBudgetOverrides,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
24pub enum ContentSource {
25 Dataset {
26 asset: AssetSelector,
27 version: Option<String>,
28 limit: Option<u64>,
29 },
30 Preview {
31 asset: AssetSelector,
32 version: Option<String>,
33 limit: u64,
34 },
35 Datafile {
36 asset: AssetSelector,
37 version: Option<String>,
38 #[serde(default = "yes")]
39 decrypt: bool,
40 #[serde(default = "yes")]
41 decompress: bool,
42 },
43 Query {
44 inputs: Vec<QueryInput>,
45 #[serde(default)]
46 views: std::collections::BTreeMap<String, String>,
47 sql: String,
48 },
49}
50fn yes() -> bool {
51 true
52}
53#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
54#[serde(deny_unknown_fields)]
55pub struct QueryInput {
56 pub alias: String,
57 pub asset: AssetSelector,
58 pub version: Option<String>,
59}
60
61pub use ahri_tre_types::ContentRepresentation;
62
63#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
64#[serde(deny_unknown_fields)]
65pub struct ContentDescriptor {
66 pub framing_version: u32,
67 pub request_id: RequestId,
68 pub transfer_id: uuid::Uuid,
69 pub admission_id: uuid::Uuid,
70 pub inputs: Vec<ObjectRef>,
71 pub representation: ContentRepresentation,
72 pub disclosure: String,
73 pub budgets: DisclosureBudgets,
74 pub deadline: chrono::DateTime<chrono::Utc>,
75}
76
77#[derive(Debug, Clone, Serialize, Deserialize)]
78#[serde(deny_unknown_fields)]
79pub struct ContentClaim {
80 pub transfer_id: uuid::Uuid,
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize)]
84#[serde(deny_unknown_fields)]
85pub struct ContentLimits {
86 pub defaults: DisclosureBudgets,
87 pub ceilings: DisclosureBudgets,
88 pub frame_bytes: usize,
89 pub arrow_message_bytes: usize,
90 pub active_transfers: usize,
91 pub idle_seconds: u32,
92}
93
94impl ContentRequest {
95 pub fn from_protocol(
98 envelope: &crate::ProtocolRequestEnvelope,
99 ) -> Result<Option<Self>, &'static str> {
100 use crate::{
101 datafile::DataFileSelector,
102 dataset::{DatasetDataFormat, DatasetSelector},
103 request::ProtocolRequest,
104 };
105 fn dataset(selector: &DatasetSelector) -> AssetSelector {
106 match selector {
107 DatasetSelector::Id { dataset, study } => AssetSelector::Id {
108 asset: *dataset,
109 study: study.clone(),
110 asset_type: Some(ahri_tre_types::AssetType::Dataset),
111 },
112 DatasetSelector::Name { study, name } => AssetSelector::Name {
113 study: study.clone(),
114 name: name.clone(),
115 asset_type: Some(ahri_tre_types::AssetType::Dataset),
116 },
117 }
118 }
119 fn format(value: DatasetDataFormat, compress: bool) -> ContentRepresentation {
120 match value {
121 DatasetDataFormat::Arrow => ContentRepresentation::ArrowIpc,
122 DatasetDataFormat::Parquet => ContentRepresentation::Parquet,
123 DatasetDataFormat::Csv if compress => ContentRepresentation::CsvGzip,
124 DatasetDataFormat::Csv => ContentRepresentation::Csv,
125 DatasetDataFormat::Json => ContentRepresentation::Json,
126 }
127 }
128 let invalid_compression = match &envelope.request {
129 ProtocolRequest::DatasetData(request) => {
130 request.compress && request.format != DatasetDataFormat::Csv
131 }
132 ProtocolRequest::DatasetExport(request) => {
133 request.compress && request.format != DatasetDataFormat::Csv
134 }
135 _ => false,
136 };
137 if invalid_compression {
138 return Err("Dataset compression requires CSV");
139 }
140 let (session, source, representation) = match &envelope.request {
141 ProtocolRequest::DatasetData(request) => (
142 request.session.clone(),
143 ContentSource::Dataset {
144 asset: dataset(&request.dataset),
145 version: None,
146 limit: request.limit.map(|n| n as u64),
147 },
148 format(request.format, request.compress),
149 ),
150 ProtocolRequest::DatasetExport(request) => (
151 request.session.clone(),
152 ContentSource::Dataset {
153 asset: dataset(&request.dataset),
154 version: None,
155 limit: request.limit.map(|n| n as u64),
156 },
157 format(request.format, request.compress),
158 ),
159 ProtocolRequest::DatasetPreview(request) => (
160 request.session.clone(),
161 ContentSource::Preview {
162 asset: dataset(&request.dataset),
163 version: None,
164 limit: request.limit as u64,
165 },
166 ContentRepresentation::Json,
167 ),
168 ProtocolRequest::DataFileExport(request) => {
169 let asset = match &request.datafile {
170 DataFileSelector::Id { datafile, study } => AssetSelector::Id {
171 asset: *datafile,
172 study: study.clone(),
173 asset_type: Some(ahri_tre_types::AssetType::File),
174 },
175 DataFileSelector::Name { study, name } => AssetSelector::Name {
176 study: study.clone(),
177 name: name.clone(),
178 asset_type: Some(ahri_tre_types::AssetType::File),
179 },
180 };
181 (
182 request.session.clone(),
183 ContentSource::Datafile {
184 asset,
185 version: request.version.clone(),
186 decrypt: true,
187 decompress: true,
188 },
189 if request.compress {
190 ContentRepresentation::FileZstd
191 } else {
192 ContentRepresentation::File
193 },
194 )
195 }
196 _ => return Ok(None),
197 };
198 Ok(Some(Self {
199 protocol_version: envelope.protocol_version.clone(),
200 request_id: envelope.request_id.unwrap_or_default(),
201 session,
202 source,
203 representation,
204 budgets: Default::default(),
205 }))
206 }
207}