Repository: Nike-Inc/koheesio

Issue:
[FEATURE] Add Support for Bearer Token in Http Module with Masking
## Is your feature request related to a problem? Please describe.
<!-- A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] -->
Currently, the `get_headers` method in the Http module does not support adding a bearer token by default. Additionally, any bearer token added to the headers is not masked in the output, which poses a security risk. The existing masking mechanism does not work effectively as we have to add "get_secret_value" to the bearer token entry in the header, rendering the 'secret' status of it null and void.

## Describe the solution you'd like
<!-- A clear and concise description of what you want to happen. -->
1. Add a default ability to include a bearer token in the headers.
2. Ensure that the bearer token is masked in any output to maintain its 'secret' status.

## Describe alternatives you've considered
<!-- A clear and concise description of any alternative solutions or features you've considered. -->
An alternative solution could be to manually add and mask the bearer token in each instance where it is required, but this is not efficient and prone to errors.

## Additional context
<!-- Add any other context or screenshots about the feature request here. -->
The current implementation of [`get_headers`](https://github.com/Nike-Inc/koheesio/blob/af7c5811879fd149a962d1993dd2152264589431/src/koheesio/steps/http.py#L184-L188) in the Http module:
```python
def get_headers(self) -> dict:
    """
    Dump headers into JSON without SecretStr masking.
    """
    return json.loads(self.model_dump_json()).get("headers")
```


Discussion:
nogitting: Why limit it to the _bearer_ type? The user should be responsible for that choice and we just incorporate it in a safe way.
dannymeijer: > Why limit it to the _bearer_ type? The user should be responsible for that choice and we just incorporate it in a safe way.

I'm open for suggestions 😇
nogitting: > I'm open for suggestions 😇

Replacing `bearer_token` with `authorization_header` will do (most of) the job. And if this attribute is set, we just update the `headers` attribute.
dannymeijer: > > I'm open for suggestions 😇
> 
> Replacing `bearer_token` with `authorization_header` will do (most of) the job. And if this attribute is set, we just update the `headers` attribute.

Hmm.. so add something like this:
```python
    authorization_header: Optional[SecretStr] = Field(
        default=None,
        description="Token for authorization",
        alias="token",
        repr=False,
    )
```

I will address this in a separate PR.

Patch:
diff --git a/src/koheesio/spark/utils/common.py b/src/koheesio/spark/utils/common.py
index 78bdf31..341ce47 100644
--- a/src/koheesio/spark/utils/common.py
+++ b/src/koheesio/spark/utils/common.py
@@ -89,7 +89,7 @@ def check_if_pyspark_connect_is_supported() -> bool:
     Raises
     ------
     ImportError
-        If the required modules for Spark Connect (grpcio and protobuf) are not importable while Spark Connect is being 
+        If the required modules for Spark Connect (grpcio and protobuf) are not importable while Spark Connect is being
         accessed.
     """
     # before pyspark 3.4, connect was not supported
@@ -107,7 +107,7 @@ def check_if_pyspark_connect_is_supported() -> bool:
             raise ImportError(
                 "It looks like the required modules for Spark Connect (e.g. grpcio) are not installed. "
                 "If not, you can install them using `pip install pyspark[connect]` or `koheesio[pyspark_connect]`. "
-            )
+            ) from e
 
     return False
 
diff --git a/src/koheesio/steps/http.py b/src/koheesio/steps/http.py
index e33f2c7..d3a797c 100644
--- a/src/koheesio/steps/http.py
+++ b/src/koheesio/steps/http.py
@@ -14,7 +14,7 @@ response = (
 In the above example, the `response` variable will contain the JSON response from the HTTP request.
 """
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, Generator, List, Optional, Union
 import contextlib
 from enum import Enum
 import json
@@ -28,6 +28,7 @@ from koheesio.models import (
     SecretStr,
     field_serializer,
     field_validator,
+    model_validator,
 )
 
 __all__ = [
@@ -97,11 +98,17 @@ class HttpStep(Step, ExtraParamsMixin):
         description="API endpoint URL",
         alias="uri",
     )
-    headers: Optional[Dict[str, Union[str, SecretStr]]] = Field(
-        default_factory=dict,
+    headers: Dict[str, Union[str, SecretStr]] = Field(
+        default={"Content-Type": "application/json"},
         description="Request headers",
         alias="header",
     )
+    bearer_token: Optional[SecretStr] = Field(
+        default=None,
+        description="Bearer token for authorization",
+        alias="token",
+        repr=False,
+    )
     data: Optional[Union[Dict[str, str], str]] = Field(
         default_factory=dict, description="[Optional] Data to be sent along with the request", alias="body"
     )
@@ -156,15 +163,22 @@ class HttpStep(Step, ExtraParamsMixin):
 
         return method_value
 
-    @field_validator("headers", mode="before")
-    def encode_sensitive_headers(cls, headers: dict) -> dict:
+    @model_validator(mode="after")
+    def encode_sensitive_headers(self) -> dict:
         """
         Encode potentially sensitive data into pydantic.SecretStr class to prevent them
         being displayed as plain text in logs.
         """
-        if auth := headers.get("Authorization"):
-            headers["Authorization"] = auth if isinstance(auth, SecretStr) else SecretStr(auth)
-        return headers
+        if token := self.bearer_token:
+            _secret_token = token.get_secret_value()
+            # ensure the token is preceded with the word 'Bearer'
+            self.headers["Authorization"] = (
+                _secret_token if _secret_token.startswith("Bearer") else f"Bearer {_secret_token}"
+            )
+            del self.bearer_token
+        if auth := self.headers.get("Authorization"):
+            self.headers["Authorization"] = auth if isinstance(auth, SecretStr) else SecretStr(auth)
+        return self
 
     @field_serializer("headers", when_used="json")
     def decode_sensitive_headers(self, headers: dict) -> dict:
@@ -214,7 +228,9 @@ class HttpStep(Step, ExtraParamsMixin):
         }
 
     @contextlib.contextmanager
-    def _request(self, method: Optional[HttpMethod] = None, stream: bool = False) -> requests.Response:
+    def _request(
+        self, method: Optional[HttpMethod] = None, stream: bool = False
+    ) -> Generator[requests.Response, None, None]:
         """
         Executes the HTTP request with retry logic.
 


Test patch:
diff --git a/tests/steps/test_http.py b/tests/steps/test_http.py
index 9f33294..ad53ef9 100644
--- a/tests/steps/test_http.py
+++ b/tests/steps/test_http.py
@@ -136,15 +136,25 @@ def test_http_step_request():
             assert response.status_code == 200
 
 
-def test_get_headers():
-    # Authorization headers are being converted into SecretStr under the hood to avoid dumping any
-    # sensitive content into logs. However, when calling the `get_headers` method, the SecretStr is being
-    # converted back to string, otherwise sensitive info would have looked like '**********'.
-    actual_headers = HttpStep(
-        url=GET_ENDPOINT,
-        headers={"Authorization": "Bearer token", "Content-Type": "application/json"},
-    ).get_headers()
+@pytest.mark.parametrize(
+    "params",
+    [
+        dict(url=GET_ENDPOINT, headers={"Authorization": "Bearer token", "Content-Type": "application/json"}),
+        dict(url=GET_ENDPOINT, headers={"Content-Type": "application/json"}, bearer_token="token"),
+        dict(url=GET_ENDPOINT, headers={"Content-Type": "application/json"}, token="token"),
+    ],
+)
+def test_get_headers(params):
+    """
+    Authorization headers are being converted into SecretStr under the hood to avoid dumping any
+    sensitive content into logs. However, when calling the `get_headers` method, the SecretStr is being
+    converted back to string, otherwise sensitive info would have looked like '**********'.
+    """
+    # Arrange and Act
+    step = HttpStep(**params)
 
+    # Assert
+    actual_headers = step.get_headers()
     assert actual_headers["Authorization"] != "**********"
     assert actual_headers["Authorization"] == "Bearer token"
     assert actual_headers["Content-Type"] == "application/json"