ahri_tre_lake/
disclosure_inputs.rs

1//! Immutable, bounded input capture. The caller holds metadata admission locks
2//! and executor exclusion while this adapter snapshots the selected relations.
3use crate::{LakeError, RestrictedQueryPlan, ScratchAttempt, ScratchAttemptId};
4use ahri_tre_types::{NcName, StudyId, VersionId};
5use std::{
6    io::{Read, Write},
7    process::{Child, ChildStderr, ChildStdout, Command, Stdio},
8    sync::{
9        Arc, Mutex,
10        atomic::{AtomicBool, Ordering},
11        mpsc,
12    },
13    time::{Duration, Instant},
14};
15
16#[derive(Clone, serde::Serialize, serde::Deserialize)]
17pub struct DisclosureDatasetInput {
18    pub table_uuid: uuid::Uuid,
19    pub lake_snapshot: i64,
20    pub study: StudyId,
21    pub version: VersionId,
22    pub name: NcName,
23    pub major: i32,
24    pub minor: i32,
25    pub patch: i32,
26}
27
28impl From<ahri_tre_types::GovernedDatasetLocation> for DisclosureDatasetInput {
29    fn from(value: ahri_tre_types::GovernedDatasetLocation) -> Self {
30        Self {
31            table_uuid: value.table_uuid,
32            lake_snapshot: value.lake_snapshot,
33            study: StudyId(value.study_id),
34            version: VersionId(value.version_id),
35            name: value.name,
36            major: value.major,
37            minor: value.minor,
38            patch: value.patch,
39        }
40    }
41}
42
43/// Captures the committed immutable physical identity while the caller retains
44/// its writer reservation (or the operator holds an upgrade maintenance window).
45pub fn dataset_physical_identity(
46    connection: &duckdb::Connection,
47    metadata_schema: &str,
48    study: StudyId,
49    name: &str,
50    major: i32,
51    minor: i32,
52    patch: i32,
53) -> Result<(uuid::Uuid, i64), LakeError> {
54    let relation = crate::dataset_loading::dataset_table_relation(study, name, major, minor, patch);
55    let schema = metadata_schema.replace('"', "\"\"");
56    let sql = format!(
57        "SELECT CAST(t.table_uuid AS VARCHAR) FROM \"__ducklake_metadata_tre_lake\".\"{schema}\".ducklake_table t JOIN \"__ducklake_metadata_tre_lake\".\"{schema}\".ducklake_schema s USING(schema_id) WHERE s.schema_name=? AND t.table_name=? AND s.end_snapshot IS NULL AND t.end_snapshot IS NULL"
58    );
59    let identity: String = connection.query_row(
60        &sql,
61        [relation.schema_name.as_str(), relation.table_name.as_str()],
62        |row| row.get(0),
63    )?;
64    let snapshot = connection.query_row(
65        "SELECT max(snapshot_id) FROM ducklake_snapshots(?)",
66        [crate::LAKE_ALIAS],
67        |row| row.get(0),
68    )?;
69    Ok((
70        identity
71            .parse()
72            .map_err(|_| LakeError::QueryNotAdmissible)?,
73        snapshot,
74    ))
75}
76
77pub struct PreparedDisclosureQuery {
78    scratch: ScratchAttempt,
79    plan: RestrictedQueryPlan,
80}
81
82impl PreparedDisclosureQuery {
83    pub fn capture(
84        authority: &crate::DisclosureCaptureAuthority,
85        plan: RestrictedQueryPlan,
86        inputs: &[DisclosureDatasetInput],
87        scratch: &ScratchAttempt,
88        deadline: chrono::DateTime<chrono::Utc>,
89        executable: &std::path::Path,
90    ) -> Result<Self, LakeError> {
91        let scratch = scratch
92            .create_child(
93                ScratchAttemptId::new(&uuid::Uuid::new_v4().simple().to_string())
94                    .map_err(|_| LakeError::QueryNotAdmissible)?,
95            )
96            .map_err(|_| LakeError::QueryNotAdmissible)?;
97        Self::capture_owned(authority, plan, inputs, scratch, deadline, executable)
98    }
99
100    /// Capture inside an owned subtree of a durable Dataset writer attempt.
101    /// The caller's enclosing reservation tree verifies cleanup before admission
102    /// or release, including failures of this capability's best-effort Drop.
103    pub fn capture_owned(
104        authority: &crate::DisclosureCaptureAuthority,
105        plan: RestrictedQueryPlan,
106        inputs: &[DisclosureDatasetInput],
107        scratch: ScratchAttempt,
108        deadline: chrono::DateTime<chrono::Utc>,
109        executable: &std::path::Path,
110    ) -> Result<Self, LakeError> {
111        if inputs
112            .iter()
113            .map(|input| input.version.0)
114            .collect::<std::collections::BTreeSet<_>>()
115            != *plan.input_versions()
116        {
117            return Err(LakeError::QueryNotAdmissible);
118        }
119        crate::capture_worker::capture(authority, inputs, scratch.path(), deadline, executable)?;
120        Ok(Self { scratch, plan })
121    }
122
123    /// Internal transform output, retained in private scratch. User SQL still
124    /// executes with no Lake or PostgreSQL authority and no ambient file access.
125    pub fn materialize(
126        self,
127        executable: &std::path::Path,
128        deadline: chrono::DateTime<chrono::Utc>,
129        max_bytes: u64,
130        max_rows: u64,
131    ) -> Result<PreparedQueryOutput, LakeError> {
132        let remaining = (deadline - chrono::Utc::now()).num_seconds();
133        if remaining < 1 || remaining > i64::from(u32::MAX) || max_bytes == 0 || max_rows == 0 {
134            return Err(LakeError::QueryNotAdmissible);
135        }
136        let path = self.scratch.path().join("output.parquet");
137        let mut command = Command::new(executable);
138        parent_bound(&mut command)?;
139        let mut child = command
140            .env_clear()
141            .arg(self.scratch.path().join("inputs.duckdb"))
142            .arg(remaining.to_string())
143            .arg(max_bytes.to_string())
144            .arg(max_rows.to_string())
145            .arg(&path)
146            .stdin(Stdio::piped())
147            .stdout(Stdio::null())
148            .stderr(Stdio::null())
149            .spawn()?;
150        if child
151            .stdin
152            .take()
153            .ok_or(LakeError::QueryNotAdmissible)?
154            .write_all(self.plan.restricted_sql().as_bytes())
155            .is_err()
156        {
157            let _ = child.kill();
158            let _ = child.wait();
159            return Err(LakeError::QueryNotAdmissible);
160        }
161        loop {
162            match child.try_wait() {
163                Ok(Some(status)) if status.success() && chrono::Utc::now() < deadline => break,
164                Ok(Some(_)) => return Err(LakeError::QueryNotAdmissible),
165                Ok(None) if chrono::Utc::now() < deadline => {
166                    std::thread::sleep(Duration::from_millis(10))
167                }
168                _ => {
169                    let _ = child.kill();
170                    let _ = child.wait();
171                    return Err(LakeError::QueryNotAdmissible);
172                }
173            }
174        }
175        if std::fs::metadata(&path)?.len() > max_bytes {
176            return Err(LakeError::QueryNotAdmissible);
177        }
178        Ok(PreparedQueryOutput {
179            path,
180            _inputs: self,
181        })
182    }
183
184    /// Executable location is deployment-owned. This API never accepts a client
185    /// path, environment or connection string as worker input authority.
186    pub fn execute(
187        self,
188        executable: &std::path::Path,
189        lifetime_seconds: u32,
190        max_bytes: u64,
191        max_rows: u64,
192    ) -> Result<RestrictedQueryStream, LakeError> {
193        self.execute_representation(
194            executable,
195            ahri_tre_types::ContentRepresentation::ArrowIpc,
196            None,
197            lifetime_seconds,
198            max_bytes,
199            max_rows,
200        )
201    }
202
203    pub fn execute_representation(
204        self,
205        executable: &std::path::Path,
206        representation: ahri_tre_types::ContentRepresentation,
207        preview_limit: Option<u64>,
208        lifetime_seconds: u32,
209        max_bytes: u64,
210        max_rows: u64,
211    ) -> Result<RestrictedQueryStream, LakeError> {
212        if matches!(
213            representation,
214            ahri_tre_types::ContentRepresentation::File
215                | ahri_tre_types::ContentRepresentation::FileZstd
216        ) {
217            return Err(LakeError::QueryNotAdmissible);
218        }
219        if lifetime_seconds == 0
220            || max_bytes == 0
221            || max_bytes > i64::MAX as u64
222            || max_rows == 0
223            || max_rows > i64::MAX as u64
224        {
225            return Err(LakeError::QueryNotAdmissible);
226        }
227        let mut command = Command::new(executable);
228        parent_bound(&mut command)?;
229        let mut child = command
230            .env_clear()
231            .arg(self.scratch.path().join("inputs.duckdb"))
232            .arg(lifetime_seconds.to_string())
233            .arg(max_bytes.to_string())
234            .arg(max_rows.to_string())
235            .arg(format!("--stream:{}", representation.as_str()))
236            .args(preview_limit.map(|limit| limit.to_string()))
237            .stdin(Stdio::piped())
238            .stdout(Stdio::piped())
239            .stderr(Stdio::piped())
240            .spawn()?;
241        let sent = child
242            .stdin
243            .take()
244            .ok_or(LakeError::QueryNotAdmissible)?
245            .write_all(self.plan.restricted_sql().as_bytes());
246        if sent.is_err() {
247            let _ = child.kill();
248            let _ = child.wait();
249            return Err(LakeError::QueryNotAdmissible);
250        }
251        let stdout = child.stdout.take().ok_or(LakeError::QueryNotAdmissible)?;
252        let completion = child.stderr.take().ok_or(LakeError::QueryNotAdmissible)?;
253        let child = Arc::new(Mutex::new(child));
254        let progress = Arc::new(Mutex::new(Instant::now()));
255        let timed_out = Arc::new(AtomicBool::new(false));
256        let timeout_signal = timed_out.clone();
257        let (stop, stopped) = mpsc::channel();
258        let process = child.clone();
259        let last_progress = progress.clone();
260        let watchdog = std::thread::spawn(move || {
261            loop {
262                let remaining = Duration::from_secs(30).saturating_sub(
263                    last_progress
264                        .lock()
265                        .unwrap_or_else(|e| e.into_inner())
266                        .elapsed(),
267                );
268                match stopped.recv_timeout(remaining) {
269                    Ok(()) | Err(mpsc::RecvTimeoutError::Disconnected) => break,
270                    Err(mpsc::RecvTimeoutError::Timeout) => {
271                        if last_progress
272                            .lock()
273                            .unwrap_or_else(|e| e.into_inner())
274                            .elapsed()
275                            >= Duration::from_secs(30)
276                        {
277                            timeout_signal.store(true, Ordering::Release);
278                            let mut child = process.lock().unwrap_or_else(|e| e.into_inner());
279                            let _ = child.kill();
280                            let _ = child.wait();
281                            break;
282                        }
283                    }
284                }
285            }
286        });
287        Ok(RestrictedQueryStream {
288            child,
289            progress,
290            timed_out,
291            stop,
292            watchdog: Some(watchdog),
293            stdout,
294            completion,
295            rows: None,
296            max_rows,
297            _inputs: self,
298            complete: false,
299        })
300    }
301}
302
303pub struct PreparedQueryOutput {
304    pub(crate) path: std::path::PathBuf,
305    _inputs: PreparedDisclosureQuery,
306}
307
308/// Dropping a stream kills and reaps its worker before releasing private inputs.
309pub struct RestrictedQueryStream {
310    child: Arc<Mutex<Child>>,
311    progress: Arc<Mutex<Instant>>,
312    timed_out: Arc<AtomicBool>,
313    stop: mpsc::Sender<()>,
314    watchdog: Option<std::thread::JoinHandle<()>>,
315    stdout: ChildStdout,
316    completion: ChildStderr,
317    rows: Option<u64>,
318    max_rows: u64,
319    _inputs: PreparedDisclosureQuery,
320    complete: bool,
321}
322impl RestrictedQueryStream {
323    /// Known only after successful worker exit and complete output consumption.
324    pub fn rows(&self) -> Option<u64> {
325        self.rows
326    }
327}
328impl Read for RestrictedQueryStream {
329    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
330        if buffer.is_empty() {
331            return Ok(0);
332        }
333        if self.timed_out.load(Ordering::Acquire) {
334            return Err(std::io::Error::other("disclosure idle timeout"));
335        }
336        let read = self.stdout.read(buffer)?;
337        if self.timed_out.load(Ordering::Acquire) {
338            return Err(std::io::Error::other("disclosure idle timeout"));
339        }
340        if read > 0 {
341            *self.progress.lock().unwrap_or_else(|e| e.into_inner()) = Instant::now();
342        }
343        if read == 0 && !self.complete {
344            if !self
345                .child
346                .lock()
347                .unwrap_or_else(|e| e.into_inner())
348                .wait()?
349                .success()
350            {
351                return Err(std::io::Error::other("restricted query did not complete"));
352            }
353            let mut completion = String::new();
354            self.completion
355                .by_ref()
356                .take(32)
357                .read_to_string(&mut completion)?;
358            let rows = completion
359                .trim()
360                .parse::<u64>()
361                .ok()
362                .filter(|rows| *rows <= self.max_rows)
363                .ok_or_else(|| std::io::Error::other("restricted query completion invalid"))?;
364            if completion != format!("{rows}\n") {
365                return Err(std::io::Error::other("restricted query completion invalid"));
366            }
367            self.rows = Some(rows);
368            self.complete = true;
369        }
370        Ok(read)
371    }
372}
373impl Drop for RestrictedQueryStream {
374    fn drop(&mut self) {
375        if !self.complete {
376            let mut child = self.child.lock().unwrap_or_else(|e| e.into_inner());
377            let _ = child.kill();
378            let _ = child.wait();
379        }
380        let _ = self.stop.send(());
381        if let Some(watchdog) = self.watchdog.take() {
382            let _ = watchdog.join();
383        }
384    }
385}
386
387/// Install the death signal before exec and verify the original parent after
388/// registration. Only async-signal-safe syscalls run in the forked child.
389pub(crate) fn parent_bound(command: &mut Command) -> std::io::Result<()> {
390    #[cfg(target_os = "linux")]
391    {
392        use std::os::unix::process::CommandExt;
393        // SAFETY: getpid takes no pointers and observes the spawning process.
394        let parent = unsafe { libc::getpid() };
395        // SAFETY: this closure only invokes async-signal-safe Linux syscalls;
396        // it does not acquire locks or allocate memory before exec.
397        unsafe {
398            command.pre_exec(move || {
399                if libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGKILL) != 0 {
400                    return Err(std::io::Error::last_os_error());
401                }
402                if libc::getppid() != parent {
403                    libc::_exit(1);
404                }
405                Ok(())
406            });
407        }
408        Ok(())
409    }
410    #[cfg(not(target_os = "linux"))]
411    Err(std::io::Error::other("worker supervision unavailable"))
412}
ahri_tre_lake/disclosure_inputs.rs

ahri_tre_lake/
disclosure_inputs.rs
