Merge pull request #7 from marklogic/feature/495-renew-token

DEVEXP-495 Token is now renewed automatically.

Author: rjrudin

CONTRIBUTING.md

@@ -39,7 +39,7 @@ ensuring you've done that, start a new Python shell by running `python`.
 
 You can then import the client via:
 
-    from marklogic.client import Client
+    from marklogic import Client
 
 You can instantiate an instance of the client that communicates with this project's test application via:
 

docs/getting-started.md

@@ -19,12 +19,12 @@ as you'd use either the `Session` class or the `requests` API.
 
 A `Client` instance can be constructed either by providing a base URL for all requests along with authentication:
 
-    from marklogic.client import Client
+    from marklogic import Client
     client = Client('http://localhost:8030', digest=('username', 'password'))
 
 Or via separate arguments for each of the parts of a base URL:
 
-    from marklogic.client import Client
+    from marklogic import Client
     client = Client(host='localhost', port='8030', digest=('username', 'password'))
 
 After constructing a `Client` instance, each of the methods in the `requests` API for sending an HTTP request can be 
@@ -43,14 +43,14 @@ Because the `Client` class extends the `Sessions` class, it can be used as a con
 
 The `Client` constructor includes a `digest` argument as a convenience for using digest authentication:
 
-    from marklogic.client import Client
+    from marklogic import Client
     client = Client('http://localhost:8030', digest=('username', 'password'))
 
 An `auth` argument is also available for using any authentication strategy that can be configured
 [via the requests `auth` argument](https://requests.readthedocs.io/en/latest/user/advanced/#custom-authentication). For 
 example, just like with `requests`, a tuple can be passed to the `auth` argument to use basic authentication:
 
-    from marklogic.client import Client
+    from marklogic import Client
     client = Client('http://localhost:8030', auth=('username', 'password'))
 
 ### MarkLogic Cloud Authentication
@@ -59,26 +59,35 @@ When connecting to a [MarkLogic Cloud instance](https://developer.marklogic.com/
 the `cloud_api_key` and `base_path` arguments. You only need to specify a `host` as well, as port 443 and HTTPS will be
 used by default. For example:
 
-    from marklogic.client import Client
+    from marklogic import Client
     client = Client(host='example.marklogic.cloud', cloud_api_key='some-key-value', base_path='/ml/example/manage')
 
 You may still use a full base URL if you wish:
 
-    from marklogic.client import Client
+    from marklogic import Client
     client = Client('https://example.marklogic.cloud', cloud_api_key='some-key-value', base_path='/ml/example/manage')
 
- 
+MarkLogic Cloud uses an access token for authentication; the access token is generated using the API key value. In some 
+scenarios, you may wish to set the token expiration time to a value other than the default used by MarkLogic Cloud. To 
+do so, set the `cloud_token_duration` argument to a number greater than zero that defines the token duration in 
+minutes:
+
+    from marklogic import Client
+    # Sets a token duration of 10 minutes.
+    client = Client(host='example.marklogic.cloud', cloud_api_key='some-key-value', base_path='/ml/example/manage', 
+        cloud_token_duration=10)
+
 ## SSL 
 
 Configuring SSL connections is the same as 
 [documented for the `requests` library](https://requests.readthedocs.io/en/latest/user/advanced/#ssl-cert-verification). 
 As a convience, the `Client` constructor includes a `verify` argument so that it does not need to be configured on the 
 `Client` instance after it's been constructed nor on every request:
 
-    from marklogic.client import Client
+    from marklogic import Client
     client = Client('https://localhost:8030', digest=('username', 'password'), verify='/path/to/cert.pem')
 
 When specifying the base URL via separate arguments, the `scheme` argument can be set for HTTPS connections:
 
-    from marklogic.client import Client
+    from marklogic import Client
     client = Client(host='localhost', port='8030', scheme='https', digest=('username', 'password'), verify=False)

marklogic/__init__.py

@@ -0,0 +1 @@
+from marklogic.client import Client

marklogic/client.py

@@ -18,6 +18,7 @@ def __init__(
         username: str = None,
         password: str = None,
         cloud_api_key: str = None,
+        cloud_token_duration: int = 0,
     ):
         super(Client, self).__init__()
         self.verify = verify
@@ -35,7 +36,9 @@ def __init__(
         elif digest:
             self.auth = HTTPDigestAuth(digest[0], digest[1])
         elif cloud_api_key:
-            self.auth = MarkLogicCloudAuth(self.base_url, cloud_api_key, self.verify)
+            self.auth = MarkLogicCloudAuth(
+                self, self.base_url, cloud_api_key, cloud_token_duration
+            )
         else:
             self.auth = HTTPDigestAuth(username, password)
 

marklogic/cloud_auth.py

@@ -1,22 +1,60 @@
-from urllib.parse import urljoin
-
+import logging
 import requests
+from requests import Response, Session, Request
 from requests.auth import AuthBase
+from urllib.parse import urljoin
 
-# See https://requests.readthedocs.io/en/latest/user/advanced/#custom-authentication
+logger = logging.getLogger(__name__)
 
 
 class MarkLogicCloudAuth(AuthBase):
-    def __init__(self, base_url: str, api_key: str, verify):
+    """
+    Handles authenticating with MarkLogic Cloud.
+    See https://requests.readthedocs.io/en/latest/user/advanced/#custom-authentication
+    for more information on custom authentication classes in requests.
+
+    Requires an instance of Session so that when a 401 is received on a request to
+    MarkLogic - which may indicate that the token has expired - a new token can be
+    generated and the original request can be resent using the same Session that
+    initially sent it.
+    """
+
+    def __init__(
+        self,
+        session: Session,
+        base_url: str,
+        api_key: str,
+        cloud_token_duration: int = 0,
+    ):
+        self._session = session
         self._base_url = base_url
-        self._verify = verify
-        self._generate_token(api_key)
+        self._api_key = api_key
+        self._cloud_token_duration = cloud_token_duration
+        self._generate_token()
+
+        # See https://docs.python-requests.org/en/latest/user/advanced/#event-hooks for
+        # more information on requests hooks.
+        self._session.hooks["response"].append(self._renew_token_if_necessary)
+
+        # Used for keeping track of whether a request has been resent with a new token
+        # after receiving a 401; avoids an infinite loop of retrying requests.
+        self.resent_request_on_401 = False
+
+    def __call__(self, request: Request):
+        # Invoked via the requests authentication framework.
+        self._add_authorization_header(request)
+        return request
+
+    def _generate_token(self):
+        params = {}
+        if self._cloud_token_duration > 0:
+            params["duration"] = self._cloud_token_duration
 
-    def _generate_token(self, api_key: str):
         response = requests.post(
             urljoin(self._base_url, "/token"),
-            data={"grant_type": "apikey", "key": api_key},
-            verify=self._verify,
+            data={"grant_type": "apikey", "key": self._api_key},
+            verify=self._session.verify,
+            params=params,
         )
 
         if response.status_code != 200:
@@ -26,6 +64,14 @@ def _generate_token(self, api_key: str):
 
         self._access_token = response.json()["access_token"]
 
-    def __call__(self, r):
-        r.headers["Authorization"] = f"Bearer {self._access_token}"
-        return r
+    def _renew_token_if_necessary(self, response: Response, *args, **kwargs):
+        if response.status_code == 401 and not self.resent_request_on_401:
+            logger.debug("Received 401; will generate new token and try request again")
+            self.resent_request_on_401 = True
+            self._generate_token()
+            self._add_authorization_header(response.request)
+            return self._session.send(response.request, *args, **kwargs)
+        self.resent_request_on_401 = False
+
+    def _add_authorization_header(self, request: Request) -> None:
+        request.headers["Authorization"] = f"Bearer {self._access_token}"

pyproject.toml

@@ -22,3 +22,8 @@ black = "^23.3.0"
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.pytest.ini_options]
+# Enables live logging; see https://docs.pytest.org/en/latest/how-to/logging.html#live-logs 
+log_cli = 1
+log_cli_level = "DEBUG"

tests/conftest.py

@@ -1,5 +1,5 @@
 import pytest
-from marklogic.client import Client
+from marklogic import Client
 
 
 @pytest.fixture

tests/test_cloud.py

@@ -1,6 +1,8 @@
+import logging
 import pytest
+import time
 
-from marklogic.client import Client
+from marklogic import Client
 
 """
 This module is intended for manual testing where the cloud_config fixture
@@ -65,6 +67,29 @@ def test_invalid_api_key(cloud_config):
     )
 
 
+@pytest.mark.skip(
+    "Skipped since it takes over a minute to run; comment this out to run it."
+)
+def test_renew_token(cloud_config):
+    if cloud_config["key"] == "changeme":
+        return
+
+    client = Client(
+        host=cloud_config["host"],
+        cloud_api_key=cloud_config["key"],
+        cloud_token_duration=1,
+        base_path=DEFAULT_BASE_PATH,
+    )
+
+    _verify_client_works(client)
+
+    logging.info("Sleeping to ensure the token will have expired on the next call")
+    time.sleep(61)
+
+    # First call should fail, resulting in a new token being generated.
+    _verify_client_works(client)
+
+
 def _new_client(cloud_config, base_path: str) -> Client:
     return Client(
         host=cloud_config["host"],