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/lib.rs

@@ -64,6 +64,48 @@ pub enum QueryError {
     StoreError { e: String },
 }
 
+#[derive(Debug, Error)]
+#[non_exhaustive]
+pub enum S3SelectPolicyError {
+    #[error("Unsupported S3 Select SQL structure: {message}")]
+    UnsupportedSqlStructure { message: String },
+
+    #[error("S3 Select query concurrency limit reached")]
+    QueryConcurrencyLimit,
+
+    #[error("S3 Select query exceeded the {seconds}-second execution limit")]
+    QueryTimeout { seconds: u64 },
+}
+
+impl S3SelectPolicyError {
+    fn from_error<'a>(mut err: &'a (dyn std::error::Error + 'static)) -> Option<&'a Self> {
+        for _ in 0..16 {
+            if let Some(policy_error) = err.downcast_ref::<Self>() {
+                return Some(policy_error);
+            }
+            err = err.source()?;
+        }
+        None
+    }
+}
+
+impl QueryError {
+    pub fn s3_select_policy_error(&self) -> Option<&S3SelectPolicyError> {
+        match self {
+            Self::Datafusion { source } => S3SelectPolicyError::from_error(source.as_ref()),
+            _ => None,
+        }
+    }
+}
+
+impl From<S3SelectPolicyError> for QueryError {
+    fn from(value: S3SelectPolicyError) -> Self {
+        Self::Datafusion {
+            source: Box::new(DataFusionError::External(Box::new(value))),
+        }
+    }
+}
+
 impl From<DataFusionError> for QueryError {
     fn from(value: DataFusionError) -> Self {
         match value {
@@ -116,9 +158,23 @@ mod tests {
         };
         assert_eq!(err.to_string(), "Multi-statement not allow, found num:2, sql:SELECT 1; SELECT 2;");
 
+        let err = S3SelectPolicyError::UnsupportedSqlStructure {
+            message: "JOIN is not supported".to_string(),
+        };
+        assert_eq!(err.to_string(), "Unsupported S3 Select SQL structure: JOIN is not supported");
+
         let err = QueryError::Cancel;
         assert_eq!(err.to_string(), "The query has been canceled");
 
+        assert_eq!(
+            S3SelectPolicyError::QueryConcurrencyLimit.to_string(),
+            "S3 Select query concurrency limit reached"
+        );
+        assert_eq!(
+            S3SelectPolicyError::QueryTimeout { seconds: 300 }.to_string(),
+            "S3 Select query exceeded the 300-second execution limit"
+        );
+
         let err = QueryError::FunctionNotExists {
             name: "my_func".to_string(),
         };
@@ -143,6 +199,35 @@ mod tests {
         }
     }
 
+    #[test]
+    fn query_error_variants_remain_source_compatible() {
+        fn exhaustive_match(err: QueryError) {
+            match err {
+                QueryError::Datafusion { .. }
+                | QueryError::NotImplemented { .. }
+                | QueryError::MultiStatement { .. }
+                | QueryError::BuildQueryDispatcher { .. }
+                | QueryError::Cancel
+                | QueryError::Parser { .. }
+                | QueryError::FunctionNotExists { .. }
+                | QueryError::FunctionExists { .. }
+                | QueryError::StoreError { .. } => {}
+            }
+        }
+
+        exhaustive_match(QueryError::Cancel);
+    }
+
+    #[test]
+    fn policy_error_is_recoverable_from_query_error() {
+        let err: QueryError = S3SelectPolicyError::QueryTimeout { seconds: 300 }.into();
+
+        assert!(matches!(
+            err.s3_select_policy_error(),
+            Some(S3SelectPolicyError::QueryTimeout { seconds: 300 })
+        ));
+    }
+
     #[test]
     fn test_query_error_from_parser_error() {
         let parser_error = ParserError::ParserError("syntax error".to_string());