ahri_tre_app/service/

acquisition.rs

use super::*;
use ahri_tre_protocol::{acquisition::AcquisitionRequest, request::ProtocolRequest};
use ahri_tre_runtime::acquisition::{AcquiredSource, AcquisitionAdmission, AcquisitionInput};

pub struct AcquisitionAuthorization<'a> {
    pub shared: Option<&'a ahri_tre_config::SharedSource>,
    pub identity: &'a ahri_tre_runtime::AuthenticationIdentityBinding,
    pub deadline: std::time::Instant,
    pub cancelled: std::sync::Arc<std::sync::atomic::AtomicBool>,
}

impl AppService {
    /// Resolve and freeze the destination independently of external source access.
    /// A shared-source name grants only its explicitly configured endpoint/Study.
    pub async fn authorize_acquisition(
        &self,
        session: &DataStoreSession,
        request: &mut AcquisitionRequest,
        authority: AcquisitionAuthorization<'_>,
    ) -> Result<(), AppError> {
        let invalid =
            || AppError::Validation("Source acquisition is unavailable or unauthorized".into());
        let _deadline = session
            .dataset_admission_repository()
            .ok_or_else(invalid)?
            .upload_deadline(authority.deadline, authority.cancelled)
            .map_err(|_| invalid())?;
        if !request.valid() {
            return Err(invalid());
        }
        let (explicit, classification) = match &request.request.request {
            ProtocolRequest::IngestFile(value) => (value.study.clone(), &value.classification),
            ProtocolRequest::IngestDatasetFile(value) => {
                (Some(value.study.clone()), &value.classification)
            }
            ProtocolRequest::IngestRedcap(value) => {
                (Some(value.study.clone()), &value.classification)
            }
            ProtocolRequest::IngestSql(value)
                if !matches!(
                    value.source,
                    ahri_tre_protocol::ingest::SqlSource::Duckdb { .. }
                        | ahri_tre_protocol::ingest::SqlSource::Sqlite { .. }
                ) =>
            {
                (Some(value.study.clone()), &value.classification)
            }
            _ => return Err(invalid()),
        };
        let study = self
            .resolve_session_study(
                self.catalogue_datastore_id()?.as_uuid(),
                explicit,
                session.current_study().map(|study| study.0),
            )
            .await
            .map_err(|_| invalid())?;
        let repository = session.dataset_admission_repository().ok_or_else(invalid)?;
        repository
            .authorize_acquisition_destination(
                ahri_tre_types::StudyId(study.id.as_uuid()),
                classification,
            )
            .map_err(|_| invalid())?;
        if let ProtocolRequest::IngestDatasetFile(value) = &request.request.request {
            let domain = self.catalogue_domain_selector(value.domain.clone())?;
            self.get_domain(GetDomainRequest { domain })
                .await?
                .ok_or_else(invalid)?;
        }
        if let ProtocolRequest::IngestSql(value) = &mut request.request.request {
            let domain = self
                .get_domain(GetDomainRequest {
                    domain: self.catalogue_domain_selector(value.domain.clone())?,
                })
                .await?
                .ok_or_else(invalid)?;
            self.ensure_study_domain(
                ahri_tre_types::StudyId(study.id.as_uuid()),
                domain.domain_id,
                "SQL acquisition",
            )
            .await?;
            value.domain = ahri_tre_protocol::domain::DomainSelector::Id {
                domain: crate::projections::domain_ref(
                    self.catalogue_datastore_id()?,
                    domain.domain_id,
                ),
            };
            DatasetOutputIntent::parse(
                value.materialization.replace,
                value.materialization.new_version.as_deref(),
            )?;
            ahri_tre_lake::SqlSourcePlan::parse(&value.sql).map_err(|_| invalid())?;
        }
        if let ProtocolRequest::IngestRedcap(value) = &mut request.request.request {
            let domain = self
                .get_domain(GetDomainRequest {
                    domain: self.catalogue_domain_selector(value.domain.clone())?,
                })
                .await?
                .ok_or_else(invalid)?;
            self.ensure_study_domain(
                ahri_tre_types::StudyId(study.id.as_uuid()),
                domain.domain_id,
                "REDCap acquisition",
            )
            .await?;
            value.domain = ahri_tre_protocol::domain::DomainSelector::Id {
                domain: crate::projections::domain_ref(
                    self.catalogue_datastore_id()?,
                    domain.domain_id,
                ),
            };
        }
        match (&request.shared_source, authority.shared) {
            (None, None) => {}
            (Some(_), Some(source))
                if request.path == ahri_tre_types::AcquisitionPath::Trusted
                    && source.endpoint == request.endpoint.as_str()
                    && source.authentication == request.authentication
                    && matches!(
                        source.authentication,
                        ahri_tre_types::SourceAuthentication::Basic
                            | ahri_tre_types::SourceAuthentication::Bearer
                            | ahri_tre_types::SourceAuthentication::PostgreSql
                            | ahri_tre_types::SourceAuthentication::MsSql
                    )
                    && source.grants.iter().any(|grant| {
                        grant.issuer.as_str() == authority.identity.issuer()
                            && grant.subject.as_str() == authority.identity.subject()
                            && grant.study == study.id.as_uuid()
                    }) => {}
            _ => return Err(invalid()),
        }
        let selector = ahri_tre_protocol::study::StudySelector::Id { study };
        match &mut request.request.request {
            ProtocolRequest::IngestFile(value) => value.study = Some(selector),
            ProtocolRequest::IngestDatasetFile(value) => value.study = selector,
            ProtocolRequest::IngestRedcap(value) => value.study = selector,
            ProtocolRequest::IngestSql(value) => value.study = selector,
            _ => unreachable!(),
        }
        Ok(())
    }

    pub async fn ingest_acquisition(
        &self,
        session: &mut DataStoreSession,
        status: ahri_tre_protocol::session::SessionStatusPayload,
        admission: &AcquisitionAdmission,
        input: AcquisitionInput,
        outbound: &ahri_tre_config::OutboundPolicy,
    ) -> Result<serde_json::Value, AppError> {
        let invalid =
            || AppError::Validation("Source acquisition is unavailable or rejected".into());
        admission.check().map_err(|_| invalid())?;
        let (source, metadata): (Box<dyn std::io::Read + Send>, _) = match input {
            AcquisitionInput::Trusted(credentials) => {
                let scratch = session.shared_scratch_attempt().ok_or_else(invalid)?;
                let request = admission.request.clone();
                let policy = outbound.clone();
                let budgets = admission.budgets;
                let deadline = admission.deadline;
                let cancelled = admission.cancelled.clone();
                let acquired = run_blocking_app_work(move || {
                    AcquiredSource::fetch(
                        &request,
                        credentials,
                        &policy,
                        budgets,
                        deadline,
                        cancelled,
                        scratch.path(),
                    )
                    .map_err(|_| AppError::Validation("Source acquisition failed".into()))
                })
                .await?;
                (Box::new(acquired.file), acquired.metadata)
            }
            AcquisitionInput::Client { metadata, source } => (source, *metadata),
        };
        if metadata.upload.filename != admission.request.filename
            || metadata.upload.media_type != admission.request.media_type
        {
            return Err(invalid());
        }
        let mut request = admission.request.request.clone();
        let redcap = matches!(&request.request, ProtocolRequest::IngestRedcap(_));
        match (&mut request.request, metadata.redcap) {
            (ProtocolRequest::IngestRedcap(value), Some(roles)) => {
                value.source = roles.into();
                value.materialization.response_encoding = None;
                value.materialization.max_attempts = None;
                value.materialization.timeout_secs = None;
            }
            (ProtocolRequest::IngestRedcap(_), None) | (_, Some(_)) => return Err(invalid()),
            (_, None) => {}
        }
        let upload = ahri_tre_protocol::upload::UploadRequest {
            request,
            metadata: metadata.upload,
            budgets: admission.request.budgets.clone(),
        };
        let input = ahri_tre_runtime::upload::UploadInput {
            source,
            budgets: admission.budgets,
            deadline: admission.deadline,
            cancelled: admission.cancelled.clone(),
        };
        if redcap {
            self.ingest_redcap_upload(session, status, upload, input, Some(admission.request.path))
                .await
        } else {
            self.ingest_upload_with_acquisition(
                session,
                status,
                upload,
                input,
                Some(admission.request.path),
            )
            .await
        }
    }
}