fix(s3select): enforce query and resource limits (#5028)

* fix(s3select): enforce query and resource limits

* fix(s3select): close query resource limit gaps

* fix(s3select): preserve timeout and stream invariants

* fix(s3select): enforce staged query limits

* fix(s3select): preserve policy error compatibility

* fix(s3select): bound error source traversal

crates/s3select-api/src/query/execution.rs

@@ -30,7 +30,7 @@ use crate::{QueryError, QueryResult};
 
 use super::Query;
 use super::logical_planner::Plan;
-use super::session::SessionCtx;
+use super::session::{QueryExecutionTracker, SessionCtx};
 
 pub struct PhaseTimer {
     phase_name: &'static str,
@@ -172,6 +172,7 @@ pub struct QueryStateMachine {
     pub session: SessionCtx,
     pub query: Query,
 
+    query_tracker: Option<QueryExecutionTracker>,
     state: RwLock<QueryState>,
     start: Instant,
 }
@@ -195,11 +196,31 @@ impl QueryStateMachine {
         Self {
             session,
             query,
+            query_tracker: None,
             state: RwLock::new(QueryState::ACCEPTING),
             start: Instant::now(),
         }
     }
 
+    pub fn begin_tracked(query: Query, session: SessionCtx, query_tracker: QueryExecutionTracker) -> QueryResult<Self> {
+        if !session.is_bound_to(&query_tracker) {
+            return Err(QueryError::Cancel);
+        }
+        let mut state_machine = Self::begin(query, session);
+        state_machine.query_tracker = Some(query_tracker);
+        Ok(state_machine)
+    }
+
+    pub fn query_tracker(&self) -> Option<&QueryExecutionTracker> {
+        self.query_tracker.as_ref()
+    }
+
+    pub fn tracker_matches_session(&self) -> bool {
+        self.query_tracker
+            .as_ref()
+            .is_some_and(|query_tracker| self.session.is_bound_to(query_tracker))
+    }
+
     pub fn begin_analyze(&self) {
         self.record_phase_timestamp("analyze", "start");
         self.translate_to(QueryState::RUNNING(RUNNING::ANALYZING));