Add support for timeout header

Author: diksha-sf

src/datacustomcode/named_credential/types/http_request_builder.py

@@ -13,17 +13,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Union
+from typing import (
+    Dict,
+    Optional,
+    Union,
+)
 
 from datacustomcode.named_credential.types.http_method import HTTPMethod
 from datacustomcode.named_credential.types.http_request import HTTPRequest
 
+# Control header carrying the per-request callout response timeout (seconds).
+CALLOUT_RESPONSE_TIMEOUT_HEADER = "ctx-callout-response-timeout-seconds"
+
 
 class HTTPRequestBuilder:
     def __init__(self) -> None:
         self._url = ""
         self._method: Union[str, HTTPMethod] = HTTPMethod.GET
         self._headers: Dict[str, str] = {}
+        self._response_timeout_seconds: Optional[int] = None
 
     def set_url(self, url: str) -> "HTTPRequestBuilder":
         """Set the symbolic Named Credential reference.
@@ -44,12 +52,37 @@ def set_method(self, method: Union[str, HTTPMethod]) -> "HTTPRequestBuilder":
         return self
 
     def set_headers(self, headers: Dict[str, str]) -> "HTTPRequestBuilder":
+        """Set the HTTP headers.
+
+        Args:
+            headers: a dictionary of HTTP headers
+        """
         self._headers = headers
         return self
 
+    def set_response_timeout_seconds(self, seconds: int) -> "HTTPRequestBuilder":
+        """Set the per-request callout response timeout, in seconds.
+
+        The ``ctx-callout-response-timeout-seconds`` control header allows to
+        configure the timeout for the callout response. This is optional.
+        If not set, the platform will use the default timeout.
+        """
+        if isinstance(seconds, bool) or not isinstance(seconds, int) or seconds <= 0:
+            raise ValueError(
+                "response timeout must be a positive integer number of seconds, "
+                f"got {seconds!r}."
+            )
+        self._response_timeout_seconds = seconds
+        return self
+
     def build(self) -> HTTPRequest:
+        headers = dict(self._headers)
+        if self._response_timeout_seconds is not None:
+            headers[CALLOUT_RESPONSE_TIMEOUT_HEADER] = str(
+                self._response_timeout_seconds
+            )
         return HTTPRequest(
             url=self._url,
             method=self._method,  # type: ignore[arg-type]
-            headers=self._headers,
+            headers=headers,
         )

src/datacustomcode/templates/function/example/chunking_with_external_callout/README.md

@@ -16,6 +16,7 @@ request = (
     .set_url(CALLOUT_URL)
     .set_method(HTTPMethod.POST)
     .set_headers({"Content-Type": "application/json"})
+    .set_response_timeout_seconds(60) # optional: per-callout timeout (seconds)
     .build()
 )
 # Body is sent verbatim (serialize it yourself); the response body is a raw string.

src/datacustomcode/templates/function/example/chunking_with_external_callout/entrypoint.py

@@ -108,6 +108,7 @@ def _analyze_chunk(chunk_text: str, runtime: Runtime) -> dict[str, str]:
         .set_url(CALLOUT_URL)
         .set_method(HTTPMethod.POST)
         .set_headers({"Content-Type": "application/json", "Accept": "application/json"})
+        .set_response_timeout_seconds(60)
         .build()
     )
 

src/datacustomcode/templates/script/examples/external_callout/README.md

@@ -17,6 +17,7 @@ request = (
     .set_url("callout:gemini")            # callout:<NC name>[/<path>]
     .set_method(HTTPMethod.POST)
     .set_headers({"Content-Type": "application/json"})
+    .set_response_timeout_seconds(60)     # optional: per-callout timeout (seconds)
     .build()
 )
 ```

src/datacustomcode/templates/script/examples/external_callout/entrypoint.py

@@ -26,6 +26,7 @@
 import json
 import logging
 
+from pyspark.sql import Column
 from pyspark.sql.functions import (
     array,
     col,
@@ -62,11 +63,12 @@
     .set_url(CALLOUT_URL)
     .set_method(HTTPMethod.POST)
     .set_headers({"Content-Type": "application/json", "Accept": "application/json"})
+    .set_response_timeout_seconds(60)
     .build()
 )
 
 
-def _gemini_body_col(text_col: "col") -> "col":  # noqa: F821
+def _gemini_body_col(text_col: Column) -> Column:
     """Build a per-row Gemini ``generateContent`` request body as a JSON string.
 
     Using ``to_json(struct(...))`` keeps the row text properly escaped inside the
@@ -97,7 +99,7 @@ def _summarize_on_driver(client: Client, text: str) -> str:
 
     envelope = json.loads(response.body) if response.body else {}
     try:
-        return envelope["candidates"][0]["content"]["parts"][0]["text"]
+        return str(envelope["candidates"][0]["content"]["parts"][0]["text"])
     except (KeyError, IndexError, TypeError):
         return ""
 
@@ -125,11 +127,11 @@ def main():
         col("_callout")["response"]["body"],
         "$.candidates[0].content.parts[0].text",
     )
-    
+
     df = df.select(
-        col(_ID_COLUMN).alias("id__c"),
-        col(_TEXT_COLUMN).alias("description__c"),
-        col(_KQ_ID_COLUMN).alias("kq_id__c"),
+        col("id__c").alias("id__c"),
+        col("Description__c").alias("description__c"),
+        col("kq_id__c").alias("kq_id__c"),
         summary.alias("summary__c"),
         col("_callout")["status"].alias("callout_status__c"),
         col("_callout")["response"]["status_code"].alias("callout_http_code__c"),

tests/test_named_credential.py

@@ -111,6 +111,41 @@ def test_builder_validates_on_build(self):
         with pytest.raises(ValidationError):
             HTTPRequestBuilder().set_url("").build()
 
+    def test_builder_sets_response_timeout_header(self):
+        request = (
+            HTTPRequestBuilder()
+            .set_url("callout:NC/path")
+            .set_headers({"Accept": "application/json"})
+            .set_response_timeout_seconds(60)
+            .build()
+        )
+        assert request.headers == {
+            "Accept": "application/json",
+            "ctx-callout-response-timeout-seconds": "60",
+        }
+
+    def test_builder_response_timeout_is_order_independent(self):
+        # Setting the timeout before set_headers must not be wiped out by the
+        # header map replacement — build() applies it on top either way.
+        request = (
+            HTTPRequestBuilder()
+            .set_url("callout:NC/path")
+            .set_response_timeout_seconds(60)
+            .set_headers({"Accept": "application/json"})
+            .build()
+        )
+        assert request.headers == {
+            "Accept": "application/json",
+            "ctx-callout-response-timeout-seconds": "60",
+        }
+
+    @pytest.mark.parametrize("bad", [0, -5, True, 1.5, "30", None])
+    def test_builder_rejects_invalid_response_timeout(self, bad):
+        with pytest.raises(ValueError):
+            HTTPRequestBuilder().set_url(
+                "callout:NC/path"
+            ).set_response_timeout_seconds(bad)
+
 
 class TestHTTPResponse:
     def test_response_defaults(self):