feat: Python requests

crepererum · crepererum · commit 406bac90d503 · 2025-10-30T11:29:14.000+01:00
diff --git a/guests/python/requirements.txt b/guests/python/requirements.txt
@@ -1,2 +1,7 @@
-# That's https://github.com/urllib3/urllib3/pull/3593 as a wheel.
-urllib3 @ https://github.com/crepererum/urllib3/releases/download/2.5.100/urllib3-2.5.100-py3-none-any.whl --hash=sha256:b48be99c923c989e8db2a41a21f2511d830a7b7e99bfde24ae7214baadf7daaf
+certifi==2025.10.5 --hash=sha256:0f212c2744a9bb6de0c56639a6f68afe01ecd92d91f14ae897c4fe7bbeeef0de
+charset_normalizer==3.4.4 --hash=sha256:7a32c560861a02ff789ad905a2fe94e3f840803362c84fecf1851cb4cf3dc37f
+idna==3.11 --hash=sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea
+requests==2.32.5 --hash=sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6
+
+# That's https://github.com/urllib3/urllib3/pull/3593 + https://github.com/golemcloud/urllib3/pull/2 as a wheel.
+urllib3 @ https://github.com/crepererum/urllib3/releases/download/2.5.101/urllib3-2.5.101-py3-none-any.whl --hash=sha256:2b2dcd0944a3b5d6ce7517cf068c232c1963a6702146695147454c46b4da71f1
diff --git a/guests/python/src/python_modules/mod.rs b/guests/python/src/python_modules/mod.rs
@@ -1147,6 +1147,12 @@ mod wit_world {
                     self.inner.take();
                 }
 
+                fn finish(&mut self) -> PyResult<FutureTrailers> {
+                    let body = self.inner.take().require_resource()?;
+                    let trailers = wasip2::http::types::IncomingBody::finish(body);
+                    Ok(FutureTrailers { inner: trailers })
+                }
+
                 fn stream(&self) -> PyResult<InputStream> {
                     let stream = self.inner()?.stream().to_pyres()?;
                     Ok(InputStream {
diff --git a/host/src/lib.rs b/host/src/lib.rs
@@ -4,6 +4,7 @@
 //! [DataFusion]: https://datafusion.apache.org/
 use std::{any::Any, io::Cursor, ops::DerefMut, sync::Arc};
 
+use ::http::HeaderName;
 use arrow::datatypes::DataType;
 use datafusion_common::{DataFusionError, Result as DataFusionResult};
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
@@ -20,7 +21,10 @@ use wasmtime_wasi_http::{
     HttpResult, WasiHttpCtx, WasiHttpView,
     bindings::http::types::ErrorCode as HttpErrorCode,
     body::HyperOutgoingBody,
-    types::{HostFutureIncomingResponse, OutgoingRequestConfig, default_send_request_handler},
+    types::{
+        DEFAULT_FORBIDDEN_HEADERS, HostFutureIncomingResponse, OutgoingRequestConfig,
+        default_send_request_handler,
+    },
 };
 
 use crate::{
@@ -112,9 +116,12 @@ impl WasiHttpView for WasmStateImpl {
 
     fn send_request(
         &mut self,
-        request: hyper::Request<HyperOutgoingBody>,
+        mut request: hyper::Request<HyperOutgoingBody>,
         config: OutgoingRequestConfig,
     ) -> HttpResult<HostFutureIncomingResponse> {
+        // Python `requests` sends this so we allow it but later drop it from the actual request.
+        request.headers_mut().remove(hyper::header::CONNECTION);
+
         // technically we could return an error straight away, but `urllib3` doesn't handle that super well, so we
         // create a future and validate the error in there (before actually starting the request of course)
 
@@ -134,6 +141,15 @@ impl WasiHttpView for WasmStateImpl {
 
         Ok(HostFutureIncomingResponse::pending(handle))
     }
+
+    fn is_forbidden_header(&mut self, name: &HeaderName) -> bool {
+        // Python `requests` sends this so we allow it but later drop it from the actual request.
+        if name == hyper::header::CONNECTION {
+            return false;
+        }
+
+        DEFAULT_FORBIDDEN_HEADERS.contains(name)
+    }
 }
 
 /// Pre-compiled WASM component.
diff --git a/host/tests/integration_tests/python/runtime/http.rs b/host/tests/integration_tests/python/runtime/http.rs
@@ -1,14 +1,14 @@
 use std::sync::Arc;
 
 use arrow::{
-    array::{Array, StringBuilder},
+    array::{Array, StringArray, StringBuilder},
     datatypes::{DataType, Field},
 };
 use datafusion_common::ScalarValue;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
 use datafusion_udf_wasm_host::{
     WasmPermissions, WasmScalarUdf,
-    http::{AllowCertainHttpRequests, Matcher},
+    http::{AllowCertainHttpRequests, HttpRequestValidator, Matcher},
 };
 use wasmtime_wasi_http::types::DEFAULT_FORBIDDEN_HEADERS;
 use wiremock::{Mock, MockServer, ResponseTemplate, matchers};
@@ -18,6 +18,46 @@ use crate::integration_tests::{
     test_utils::ColumnarValueExt,
 };
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_requests_simple() {
+    const CODE: &str = r#"
+import requests
+
+def perform_request(url: str) -> str:
+    return requests.get(url).text
+"#;
+
+    let server = MockServer::start().await;
+    Mock::given(matchers::any())
+        .respond_with(ResponseTemplate::new(200).set_body_string("hello world!"))
+        .expect(1)
+        .mount(&server)
+        .await;
+
+    let mut permissions = AllowCertainHttpRequests::new();
+    permissions.allow(Matcher {
+        method: http::Method::GET,
+        host: server.address().ip().to_string().into(),
+        port: server.address().port(),
+    });
+    let udf = python_udf_with_permissions(CODE, permissions).await;
+
+    let array = udf
+        .invoke_with_args(ScalarFunctionArgs {
+            args: vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(server.uri())))],
+            arg_fields: vec![Arc::new(Field::new("uri", DataType::Utf8, true))],
+            number_rows: 1,
+            return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
+        })
+        .unwrap()
+        .unwrap_array();
+
+    assert_eq!(
+        array.as_ref(),
+        &StringArray::from_iter([Some("hello world!".to_owned()),]) as &dyn Array,
+    );
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_urllib3_unguarded_fail() {
     const CODE: &str = r#"
@@ -131,6 +171,7 @@ def perform_request(url: str) -> str:
 #[tokio::test(flavor = "multi_thread")]
 async fn test_integration() {
     const CODE: &str = r#"
+import requests
 import urllib3
 
 def _headers_str_to_dict(headers: str) -> dict[str, str]:
@@ -152,14 +193,38 @@ def _headers_dict_to_str(headers: dict[str, str]) -> str:
     else:
         return headers
 
-def perform_request(method: str, url: str, headers: str | None, body: str | None) -> str:
+def test_requests(method: str, url: str, headers: str | None, body: str | None) -> str:
+    try:
+        resp = requests.request(
+            method=method,
+            url=url,
+            headers=_headers_str_to_dict(headers),
+            data=body,
+        )
+    except requests.ConnectionError as e:
+        (e,) = e.args
+        assert isinstance(e, Exception)
+        return f"ERR: {e}"
+    except Exception as e:
+        return f"ERR: {e}"
+
+    resp_status = resp.status_code
+    resp_body = f"'{resp.text}'" if resp.text else "n/a"
+    resp_headers = _headers_dict_to_str(resp.headers)
+
+    return f"OK: status={resp_status} headers={resp_headers} body={resp_body}"
+
+def test_urllib3(method: str, url: str, headers: str | None, body: str | None) -> str:
     try:
         resp = urllib3.request(
             method=method,
             url=url,
             headers=_headers_str_to_dict(headers),
             body=body,
         )
+    except urllib3.exceptions.MaxRetryError as e:
+        e = e.reason
+        return f"ERR: {e}"
     except Exception as e:
         return f"ERR: {e}"
 
@@ -169,6 +234,7 @@ def perform_request(method: str, url: str, headers: str | None, body: str | None
 
     return f"OK: status={resp_status} headers={resp_headers} body={resp_body}"
 "#;
+    const NUMBER_OF_IMPLEMENTATIONS: usize = 2;
 
     let mut cases = vec![
         TestCase {
@@ -246,16 +312,32 @@ def perform_request(method: str, url: str, headers: str | None, body: str | None
         },
         TestCase {
             base: Some("http://test.com"),
-            resp: Err("HTTPConnectionPool(host='test.com', port=80): Max retries exceeded with url: / (Caused by ProtocolError('Connection aborted.', WasiErrorCode('Request failed with wasi http error ErrorCode_HttpRequestDenied')))".to_owned()),
+            resp: Err("('Connection aborted.', WasiErrorCode('Request failed with wasi http error ErrorCode_HttpRequestDenied'))".to_owned()),
+            ..Default::default()
+        },
+        // Python `requests` sends this so we allow it but later drop it from the actual request.
+        TestCase {
+            path: "/forbidden_header/connection".to_owned(),
+            requ_headers: vec![(http::header::CONNECTION.to_string(), &["foo"])],
+            resp: Ok(TestResponse {
+                body: Some("header is filtered"),
+                ..Default::default()
+            }),
             ..Default::default()
         },
     ];
-    cases.extend(DEFAULT_FORBIDDEN_HEADERS.iter().map(|h| TestCase {
-        path: format!("/forbidden_header/{h}"),
-        requ_headers: vec![(h.to_string(), &["foo"])],
-        resp: Err("Err { value: HeaderError_Forbidden }".to_owned()),
-        ..Default::default()
-    }));
+    cases.extend(
+        DEFAULT_FORBIDDEN_HEADERS
+            .iter()
+            // Python `requests` sends this so we allow it but later drop it from the actual request.
+            .filter(|h| *h != http::header::CONNECTION)
+            .map(|h| TestCase {
+                path: format!("/forbidden_header/{h}"),
+                requ_headers: vec![(h.to_string(), &["foo"])],
+                resp: Err("Err { value: HeaderError_Forbidden }".to_owned()),
+                ..Default::default()
+            }),
+    );
 
     let server = MockServer::start().await;
     let mut permissions = AllowCertainHttpRequests::default();
@@ -267,7 +349,7 @@ def perform_request(method: str, url: str, headers: str | None, body: str | None
     let mut builder_result = StringBuilder::new();
 
     for case in &cases {
-        if let Some(mock) = case.mock(&server) {
+        if let Some(mock) = case.mock(&server, NUMBER_OF_IMPLEMENTATIONS) {
             mock.mount(&server).await;
         }
         permissions.allow(case.matcher(&server));
@@ -309,37 +391,33 @@ def perform_request(method: str, url: str, headers: str | None, body: str | None
         }
     }
 
-    let udfs = WasmScalarUdf::new(
-        python_component().await,
-        &WasmPermissions::new().with_http(permissions),
-        CODE.to_string(),
-    )
-    .await
-    .unwrap();
-    assert_eq!(udfs.len(), 1);
-    let udf = udfs.into_iter().next().unwrap();
-
-    let array = udf
-        .invoke_with_args(ScalarFunctionArgs {
-            args: vec![
-                ColumnarValue::Array(Arc::new(builder_method.finish())),
-                ColumnarValue::Array(Arc::new(builder_url.finish())),
-                ColumnarValue::Array(Arc::new(builder_headers.finish())),
-                ColumnarValue::Array(Arc::new(builder_body.finish())),
-            ],
-            arg_fields: vec![
-                Arc::new(Field::new("method", DataType::Utf8, true)),
-                Arc::new(Field::new("url", DataType::Utf8, true)),
-                Arc::new(Field::new("headers", DataType::Utf8, true)),
-                Arc::new(Field::new("body", DataType::Utf8, true)),
-            ],
-            number_rows: cases.len(),
-            return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
-        })
-        .unwrap()
-        .unwrap_array();
-
-    assert_eq!(array.as_ref(), &builder_result.finish() as &dyn Array,);
+    let args = ScalarFunctionArgs {
+        args: vec![
+            ColumnarValue::Array(Arc::new(builder_method.finish())),
+            ColumnarValue::Array(Arc::new(builder_url.finish())),
+            ColumnarValue::Array(Arc::new(builder_headers.finish())),
+            ColumnarValue::Array(Arc::new(builder_body.finish())),
+        ],
+        arg_fields: vec![
+            Arc::new(Field::new("method", DataType::Utf8, true)),
+            Arc::new(Field::new("url", DataType::Utf8, true)),
+            Arc::new(Field::new("headers", DataType::Utf8, true)),
+            Arc::new(Field::new("body", DataType::Utf8, true)),
+        ],
+        number_rows: cases.len(),
+        return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
+    };
+    let array_result = builder_result.finish();
+
+    let udfs = python_udfs_with_permissions(CODE, permissions).await;
+    assert_eq!(udfs.len(), NUMBER_OF_IMPLEMENTATIONS);
+
+    for udf in udfs {
+        println!("{}", udf.name());
+
+        let array = udf.invoke_with_args(args.clone()).unwrap().unwrap_array();
+        assert_eq!(array.as_ref(), &array_result as &dyn Array);
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -390,7 +468,7 @@ impl TestCase {
         }
     }
 
-    fn mock(&self, server: &MockServer) -> Option<Mock> {
+    fn mock(&self, server: &MockServer, hits: usize) -> Option<Mock> {
         let Self {
             base,
             method,
@@ -417,6 +495,11 @@ impl TestCase {
             ));
 
         for (k, v) in requ_headers {
+            // Python `requests` sends this so we allow it but later drop it from the actual request.
+            if k.as_str() == http::header::CONNECTION {
+                continue;
+            }
+
             builder = builder.and(matchers::headers(k.as_str(), v.to_vec()));
         }
 
@@ -430,9 +513,9 @@ impl TestCase {
             resp_template = resp_template.set_body_string(resp_body);
         }
 
-        let mock = builder
-            .respond_with(resp_template)
-            .expect(resp.is_ok() as u64);
+        let expect = if resp.is_ok() { hits as u64 } else { 0 };
+
+        let mock = builder.respond_with(resp_template).expect(expect);
         Some(mock)
     }
 }
@@ -477,3 +560,25 @@ impl wiremock::Match for NoForbiddenHeaders {
             .all(|h| !request.headers.contains_key(h))
     }
 }
+
+async fn python_udfs_with_permissions<V>(code: &'static str, permissions: V) -> Vec<WasmScalarUdf>
+where
+    V: HttpRequestValidator,
+{
+    WasmScalarUdf::new(
+        python_component().await,
+        &WasmPermissions::new().with_http(permissions),
+        code.to_owned(),
+    )
+    .await
+    .unwrap()
+}
+
+async fn python_udf_with_permissions<V>(code: &'static str, permissions: V) -> WasmScalarUdf
+where
+    V: HttpRequestValidator,
+{
+    let udfs = python_udfs_with_permissions(code, permissions).await;
+    assert_eq!(udfs.len(), 1);
+    udfs.into_iter().next().unwrap()
+}
