Skip to main content

ahri_tre_protocol/
content.rs

1//! JSON control documents for the separately framed content channel.
2use 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    /// Converts existing public content forms into path-free admission intent.
96    /// Destination handling stays with the client that owns the original request.
97    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}