feat: add support for setting x-request-id header (#606)

Author: brahmlower

censys/common/base.py

@@ -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",
@@ -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.

pyproject.toml

@@ -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"

tests/test_base.py

@@ -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