Access token and auto add api key paths (#283)

* Added support for using an access token for authorization when making requests.

-Giovanni and George

* MNT: Automatically add /api-key or /authorized as necessary

We can detect when different routes should be hit, so we can
automatically handle this for the user. This makes it so that
a user doesn't need to set the URL differently based on whether
they want to hit an authorized endpoint or not.

---------

Co-authored-by: Harrison <pleasant@menloinnovations.com>

Author: greglucas

README.md

@@ -139,22 +139,31 @@ To access some unreleased data products and quicklooks, you may
 need elevated permissions. To programmatically get that, you need
 an API Key, which can be requested from the SDC team.
 
-To use the API Key you can set environment variables and then use
-the tool as usual. Note that the api endpoints are prefixed with `/api-key`
-to request unreleased data. This will also require an update to the
-data access url. So the following should be used when programatically
-accessing the data.
+To use the API Key you can set the `IMAP_API_KEY` environment variable and then use
+the tool as usual.
 
 ```bash
-IMAP_API_KEY=<your-api-key> IMAP_DATA_ACCESS_URL=https://api.dev.imap-mission.com/api-key imap-data-access ...
+IMAP_API_KEY=<your-api-key> imap-data-access ...
 ```
 
 or with CLI flags
 
 ```bash
-imap-data-access --api-key <your-api-key> --url https://api.dev.imap-mission.com/api-key ...
+imap-data-access --api-key <your-api-key> ...
 ```
 
+### Automated use with Access token
+
+An alternative to using an API key to access protected data is using an access token provided by LASP's authentication server. LASP's authentication uses [keycloak authentication](https://www.keycloak.org/documentation).
+
+To use an access token with imap-data-access you can set the following environment variable:
+
+```dotenv
+IMAP_ACCESS_TOKEN={{Access token from LASP auth server}}
+```
+
+Any queries or downloads made with imap-data-access will now use these credentials.
+
 ## Troubleshooting
 
 ### Network issues

imap_data_access/__init__.py

@@ -58,6 +58,7 @@
     or "https://api.imap-mission.com",
     "DATA_DIR": Path(os.getenv("IMAP_DATA_DIR") or Path.cwd() / "data"),
     "API_KEY": os.getenv("IMAP_API_KEY"),
+    "ACCESS_TOKEN": os.getenv("IMAP_ACCESS_TOKEN"),
     # Create a base64 encoded string for the username and password
     # echo -n 'username:password' | base64
     "WEBPODA_TOKEN": os.getenv("IMAP_WEBPODA_TOKEN"),

imap_data_access/io.py

@@ -32,10 +32,16 @@ def _make_request(request: requests.PreparedRequest):
     when making HTTP requests and yield the response body.
     """
     logger.debug("Making request: %s", request)
+
     if imap_data_access.config["API_KEY"]:
         # Add the API key to the request headers if it exists
         request.headers["x-api-key"] = imap_data_access.config["API_KEY"]
-
+    elif imap_data_access.config["ACCESS_TOKEN"]:
+        # Add the access token to the request headers if it exists
+        # and API key does not exist
+        request.headers["Authorization"] = (
+            f"Bearer {imap_data_access.config['ACCESS_TOKEN']}"
+        )
     try:
         with requests.Session() as session:
             response = session.send(request)
@@ -50,6 +56,23 @@ def _make_request(request: requests.PreparedRequest):
         raise IMAPDataAccessError(error_msg) from e
 
 
+def _get_base_url() -> str:
+    """Get the base URL of the data access API.
+
+    Adds in the /api-key and /authorized to direct the url
+    to the proper authorized endpoints as needed.
+    """
+    url = imap_data_access.config["DATA_ACCESS_URL"]
+
+    # Only add these if someone hasn't already added the /api-key themselves.
+    if imap_data_access.config["API_KEY"] and not url.endswith("/api-key"):
+        url = f"{url}/api-key"
+    elif imap_data_access.config["ACCESS_TOKEN"] and not url.endswith("/authorized"):
+        url = f"{url}/authorized"
+
+    return url
+
+
 def download(file_path: Union[Path, str]) -> Path:
     """Download a file from the data archive.
 
@@ -78,7 +101,7 @@ def download(file_path: Union[Path, str]) -> Path:
         logger.info("The file %s already exists, skipping download", destination)
         return destination
 
-    url = f"{imap_data_access.config['DATA_ACCESS_URL']}/download/{file_path}"
+    url = f"{_get_base_url()}/download/{file_path}"
     logger.info("Downloading file %s from %s to %s", file_path, url, destination)
 
     # Create a request with the provided URL
@@ -279,7 +302,7 @@ def query(
         else:
             query_params["repointing"] = int(repointing)
 
-    url = f"{imap_data_access.config['DATA_ACCESS_URL']}/query"
+    url = f"{_get_base_url()}/query"
     request = requests.Request(method="GET", url=url, params=query_params).prepare()
 
     logger.info("Querying data archive for %s with url %s", query_params, request.url)
@@ -383,7 +406,7 @@ def reprocess(
     ):
         raise ValueError("Not a valid end date, use format 'YYYYMMDD'.")
     reprocess_params["reprocessing"] = "True"
-    url = f"{imap_data_access.config['DATA_ACCESS_URL']}/reprocess"
+    url = f"{_get_base_url()}/reprocess"
     request = requests.Request(
         method="POST", url=url, params=reprocess_params
     ).prepare()
@@ -410,7 +433,7 @@ def upload(file_path: Union[Path, str]) -> None:
         raise FileNotFoundError(file_path)
 
     # The upload name needs to be given as a path parameter
-    url = f"{imap_data_access.config['DATA_ACCESS_URL']}/upload/{file_path.name}"
+    url = f"{_get_base_url()}/upload/{file_path.name}"
     logger.info("Uploading file %s to %s", file_path, url)
 
     # We send a GET request with the filename and the server

pyproject.toml

@@ -61,7 +61,7 @@ lint.ignore = ["D104", "D203", "D213", "D413", "PLR2004"]
 
 [tool.ruff.lint.per-file-ignores]
 # S101: Use of assert detected
-"tests/*" = ["S101", "D"]
+"tests/*" = ["D", "PLR0913", "S101"]
 "*.ipynb" = ["E501"]
 "imap_data_access/io.py" = ["PLR0913", "S310"]
 "imap_data_access/webpoda.py" = ["S310"]

tests/conftest.py

@@ -15,7 +15,7 @@ def _set_global_config(monkeypatch: pytest.fixture, tmp_path: pytest.fixture):
         imap_data_access.config, "DATA_ACCESS_URL", "https://api.test.com"
     )
     # Make sure we don't leak any of this content if a user has set them locally
-    monkeypatch.setitem(imap_data_access.config, "API_KEY", "test_key")
+    monkeypatch.setitem(imap_data_access.config, "API_KEY", None)
     monkeypatch.setitem(imap_data_access.config, "WEBPODA_TOKEN", "test_token")
 
 

tests/test_config_options.py

@@ -16,6 +16,7 @@
         ("DATA_DIR", Path.cwd() / "data", str(Path("/test/path"))),
         ("DATA_ACCESS_URL", "https://api.imap-mission.com", "https://test.url"),
         ("API_KEY", None, "test-api-key"),
+        ("ACCESS_TOKEN", None, "test-access-token"),
     ],
 )
 def test_configuration_updates(config_var, default, expected):

tests/test_io.py

@@ -10,12 +10,52 @@
 import requests
 
 import imap_data_access
-from imap_data_access.io import _make_request
+from imap_data_access.io import _get_base_url, _make_request
 
 test_science_filename = "imap_swe_l1_test-description_20100101_v000.cdf"
 test_science_path = "imap/swe/l1/2010/01/" + test_science_filename
 
 
+@pytest.mark.parametrize(
+    ("url", "api_key", "access_token", "expected"),
+    [
+        # Default return base url
+        ("https://api.test.com", None, None, "https://api.test.com"),
+        # API Key appends /api-key
+        ("https://api.test.com", "test_key", None, "https://api.test.com/api-key"),
+        # API Key with already specified /api-key doesn't double add it
+        (
+            "https://api.test.com/api-key",
+            "test_key",
+            None,
+            "https://api.test.com/api-key",
+        ),
+        # Access token appends /authorized
+        ("https://api.test.com", None, "test_token", "https://api.test.com/authorized"),
+        # Access token with already specified /authorized doesn't double add it
+        (
+            "https://api.test.com/authorized",
+            None,
+            "test_token",
+            "https://api.test.com/authorized",
+        ),
+        # API Key takes precedence over access token
+        (
+            "https://api.test.com",
+            "test_key",
+            "test_token",
+            "https://api.test.com/api-key",
+        ),
+    ],
+)
+def test_base_url(url, api_key, access_token, expected, monkeypatch):
+    """Test that the base URL is set correctly based on the config."""
+    monkeypatch.setitem(imap_data_access.config, "DATA_ACCESS_URL", url)
+    monkeypatch.setitem(imap_data_access.config, "API_KEY", api_key)
+    monkeypatch.setitem(imap_data_access.config, "ACCESS_TOKEN", access_token)
+    assert _get_base_url() == expected
+
+
 def test_redirect(mock_send_request):
     """Verify that we follow a 307 redirect from newly created s3 buckets.
 
@@ -327,13 +367,23 @@ def test_upload_no_file(mock_send_request):
     "upload_file_path", ["a/b/test-file.txt", Path("a/b/test-file.txt")]
 )
 @pytest.mark.parametrize(
-    ("api_key", "expected_header"),
-    [(None, {}), ("test-api-key", {"x-api-key": "test-api-key"})],
+    ("api_key", "access_token", "expected_header"),
+    [
+        (None, None, {}),
+        ("test-api-key", None, {"x-api-key": "test-api-key"}),
+        (None, "test-access-token", {"Authorization": "Bearer test-access-token"}),
+        (
+            "test-api-key-default",
+            "test-access-token",
+            {"x-api-key": "test-api-key-default"},
+        ),
+    ],
 )
 def test_upload(
     mock_send_request,
     upload_file_path: str | Path,
     api_key: str | None,
+    access_token: str | None,
     expected_header: dict,
     monkeypatch,
 ):
@@ -346,10 +396,13 @@ def test_upload(
         The upload file path to test with
     api_key : str or None
         The API key to use for the upload
+    access_token : str or None
+        The access token to use for the upload
     expected_header : dict
         The expected header to be sent with the request
     """
     monkeypatch.setitem(imap_data_access.config, "API_KEY", api_key)
+    monkeypatch.setitem(imap_data_access.config, "ACCESS_TOKEN", access_token)
     mock_send_request.return_value.json.return_value = "https://s3-test-bucket.com"
     # Call the upload function
     file_to_upload = imap_data_access.config["DATA_DIR"] / upload_file_path
@@ -373,9 +426,14 @@ def test_upload(
     ]
 
     # First urlopen call should be to get the s3 upload url
+    auth_path = ""
+    if api_key:
+        auth_path = "/api-key"
+    elif access_token:
+        auth_path = "/authorized"
     request_sent = mock_calls[0].args[0]
     called_url = request_sent.url
-    expected_url_encoded = "https://api.test.com/upload/test-file.txt"
+    expected_url_encoded = f"https://api.test.com{auth_path}/upload/test-file.txt"
     assert called_url == expected_url_encoded
     assert request_sent.method == "GET"
     # An API key needs to be added to the header for uploads