Skip to content

Commit

Permalink
feat: add support for setting x-request-id header (#606)
Browse files Browse the repository at this point in the history
  • Loading branch information
brahmlower committed Mar 28, 2024
1 parent 2c13006 commit d5533d1
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 1 deletion.
24 changes: 24 additions & 0 deletions censys/common/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ def __init__(
self._session.proxies.update(proxies)
if cookies:
self._session.cookies.update(cookies)
self.request_id = kwargs.get("request_id")
self._session.headers.update(
{
"accept": "application/json, */8",
Expand All @@ -112,6 +113,29 @@ def __init__(
}
)

@property
def request_id(self) -> Optional[str]:
"""The x-request-id header value for API requests.
The x-request-id header is not set when the value is None.
Value is None by default
Returns:
Type[Optional[str]]: The value of the header.
"""
value = self._session.headers.get("x-request-id")
if not isinstance(value, str):
return None
return value

@request_id.setter
def request_id(self, value: Optional[str]):
if value is None:
self._session.headers.pop("x-request-id", None)
return

self._session.headers["x-request-id"] = value

@staticmethod
def _get_exception_class(_: Response) -> Type[CensysAPIException]:
"""Maps HTTP status code or ASM error code to exception.
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "censys"
version = "2.2.11"
version = "2.2.12"
description = "An easy-to-use and lightweight API wrapper for Censys APIs (censys.io)."
authors = ["Censys, Inc. <support@censys.io>"]
license = "Apache-2.0"
Expand Down
13 changes: 13 additions & 0 deletions tests/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,19 @@ def test_user_agent(self):
== requests.utils.default_user_agent() + " test"
)

def test_request_id(self):
id_value = "my-request-id"

# Test request id value is present
base = CensysAPIBase(TEST_URL, request_id=id_value)
assert base.request_id == id_value
assert base._session.headers.get("x-request-id") == id_value

# Test request id value is not present
base.request_id = None
assert base.request_id is None
assert base._session.headers.get("x-request-id") is None

@pytest.mark.filterwarnings("ignore:HTTP proxies will not be used.")
def test_proxies(self):
# Mock/actual call
Expand Down

0 comments on commit d5533d1

Please sign in to comment.