Skip to content

Commit

Permalink
Add support for bearer tokens
Browse files Browse the repository at this point in the history
  • Loading branch information
mcarans committed Mar 13, 2024
1 parent 017b72d commit 7649014
Show file tree
Hide file tree
Showing 7 changed files with 80 additions and 12 deletions.
4 changes: 2 additions & 2 deletions .github/workflows/publish.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,13 @@ jobs:
runs-on: ubuntu-latest

steps:
- uses: actions/checkout@v3
- uses: actions/checkout@v4
- name: Get history and tags for versioning to work
run: |
git fetch --prune --unshallow
git fetch --depth=1 origin +refs/tags/*:refs/tags/*
- name: Set up Python
uses: actions/setup-python@v4
uses: actions/setup-python@v5
with:
python-version: '3.x'
- name: Install dependencies
Expand Down
4 changes: 2 additions & 2 deletions .github/workflows/run-python-tests.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,9 @@ jobs:
runs-on: ubuntu-latest

steps:
- uses: actions/checkout@v3
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v4
uses: actions/setup-python@v5
with:
python-version: '3.x'
- name: Install dependencies
Expand Down
4 changes: 2 additions & 2 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -81,15 +81,15 @@ openpyxl==3.1.2
# via hdx-python-utilities (pyproject.toml)
packaging==24.0
# via pytest
petl==1.7.14
petl==1.7.15
# via frictionless
platformdirs==4.2.0
# via virtualenv
pluggy==1.4.0
# via pytest
pre-commit==3.6.2
# via hdx-python-utilities (pyproject.toml)
pydantic==2.6.3
pydantic==2.6.4
# via frictionless
pydantic-core==2.16.3
# via pydantic
Expand Down
4 changes: 4 additions & 0 deletions src/hdx/utilities/downloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ class Download(BaseDownload):
auth (Tuple[str, str]): Authorisation information in tuple form (user, pass) OR
basic_auth (str): Authorisation information in basic auth string form (Basic xxxxxxxxxxxxxxxx) OR
basic_auth_file (str): Path to file containing authorisation information in basic auth string form (Basic xxxxxxxxxxxxxxxx)
bearer_token (str): Bearer token string OR
bearer_token_file (str): Path to file containing bearer token string OR
extra_params_dict (Dict[str, str]): Extra parameters to put on end of url as a dictionary OR
extra_params_json (str): Path to JSON file containing extra parameters to put on end of url OR
extra_params_yaml (str): Path to YAML file containing extra parameters to put on end of url
Expand Down Expand Up @@ -1265,6 +1267,8 @@ def generate_downloaders(
auth (Tuple[str, str]): Authorisation information in tuple form (user, pass) OR
basic_auth (str): Authorisation information in basic auth string form (Basic xxxxxxxxxxxxxxxx) OR
basic_auth_file (str): Path to file containing authorisation information in basic auth string form (Basic xxxxxxxxxxxxxxxx)
bearer_token (str): Bearer token string OR
bearer_token_file (str): Path to file containing bearer token string OR
extra_params_dict (Dict[str, str]): Extra parameters to put on end of url as a dictionary OR
extra_params_json (str): Path to JSON file containing extra parameters to put on end of url OR
extra_params_yaml (str): Path to YAML file containing extra parameters to put on end of url
Expand Down
45 changes: 40 additions & 5 deletions src/hdx/utilities/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,9 @@ def get_session(
) -> requests.Session:
"""Set up and return Session object that is set up with retrying. Requires
either global user agent to be set or appropriate user agent parameter(s)
to be completed. If the EXTRA_PARAMS or BASIC_AUTH environment variable is
supplied, the extra_params* parameters will be ignored.
to be completed. If the EXTRA_PARAMS, BASIC_AUTH or BEARER_TOKEN
environment variable is supplied, the extra_params* parameters will be
ignored.
Args:
user_agent (Optional[str]): User agent string. HDXPythonUtilities/X.X.X- is prefixed.
Expand All @@ -48,7 +49,9 @@ def get_session(
**kwargs: See below
auth (Tuple[str, str]): Authorisation information in tuple form (user, pass) OR
basic_auth (str): Authorisation information in basic auth string form (Basic xxxxxxxxxxxxxxxx) OR
basic_auth_file (str): Path to file containing authorisation information in basic auth string form (Basic xxxxxxxxxxxxxxxx)
basic_auth_file (str): Path to file containing authorisation information in basic auth string form (Basic xxxxxxxxxxxxxxxx) OR
bearer_token (str): Bearer token string OR
bearer_token_file (str): Path to file containing bearer token string OR
extra_params_dict (Dict): Extra parameters to put on end of url as a dictionary OR
extra_params_json (str): Path to JSON file containing extra parameters to put on end of url OR
extra_params_yaml (str): Path to YAML file containing extra parameters to put on end of url
Expand Down Expand Up @@ -76,11 +79,16 @@ def get_session(
extra_params_found = False
extra_params_dict = None
basic_auth = None
bearer_token = None
if use_env:
basic_auth_env = os.getenv("BASIC_AUTH")
if basic_auth_env:
basic_auth = basic_auth_env
auths_found.append("basic_auth environment variable")
auths_found.append("BASIC_AUTH environment variable")
bearer_token_env = os.getenv("BEARER_TOKEN")
if bearer_token_env:
bearer_token = bearer_token_env
auths_found.append("BEARER_TOKEN environment variable")
extra_params = os.getenv("EXTRA_PARAMS")
if extra_params:
if "=" in extra_params:
Expand Down Expand Up @@ -133,17 +141,26 @@ def get_session(
)
if extra_params_dict:
basic_auth_param = extra_params_dict.get("basic_auth")
bearer_token_param = extra_params_dict.get("bearer_token")
if basic_auth_param:
basic_auth = basic_auth_param
auths_found.append("basic_auth parameter")
del extra_params_dict["basic_auth"]
if bearer_token_param:
bearer_token = bearer_token_param
auths_found.append("bearer_token parameter")
del extra_params_dict["bearer_token"]

s.params = extra_params_dict

basic_auth_arg = kwargs.get("basic_auth")
if basic_auth_arg:
basic_auth = basic_auth_arg
auths_found.append("basic_auth argument")
bearer_token_arg = kwargs.get("bearer_token")
if bearer_token_arg:
bearer_token = bearer_token_arg
auths_found.append("bearer_token argument")

auth = kwargs.get("auth")
if auth:
Expand All @@ -157,6 +174,15 @@ def get_session(
except OSError:
if fail_on_missing_file:
raise
bearer_token_file = kwargs.get("bearer_token_file")
if bearer_token_file:
logger.info(f"Loading bearer token from: {bearer_token_file}")
try:
bearer_token = load_text(bearer_token_file, strip=True)
auths_found.append(f"file {bearer_token_file}")
except OSError:
if fail_on_missing_file:
raise
if len(auths_found) > 1:
auths_found_str = ", ".join(auths_found)
raise SessionError(
Expand All @@ -165,7 +191,16 @@ def get_session(
if "headers" not in auths_found:
if basic_auth:
auth = basicauth_decode(basic_auth)
s.auth = auth
s.auth = auth
elif bearer_token:
s.headers.update(
{
"Accept": "application/json",
"Authorization": f"Bearer {bearer_token}",
}
)
else:
s.auth = auth

status_forcelist = kwargs.get(
"status_forcelist", (429, 500, 502, 503, 504)
Expand Down
1 change: 1 addition & 0 deletions tests/fixtures/downloader/bearertoken.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
12345
30 changes: 29 additions & 1 deletion tests/hdx/utilities/test_downloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,14 +114,29 @@ def test_get_path_for_url(
assert abspath(path) == abspath(join(downloaderfolder, filename))

def test_init(self, monkeypatch, downloaderfolder):
basicauthfile = join(downloaderfolder, "basicauth.txt")
with Download(auth=("u", "p")) as downloader:
assert downloader.session.auth == ("u", "p")
basicauth = "Basic dXNlcjpwYXNz"
with Download(basic_auth=basicauth) as downloader:
assert downloader.session.auth == ("user", "pass")
basicauthfile = join(downloaderfolder, "basicauth.txt")
with Download(basic_auth_file=basicauthfile) as downloader:
assert downloader.session.auth == ("testuser", "testpass")
bearertoken = "ABCDE"
with Download(bearer_token=bearertoken) as downloader:
assert downloader.session.headers["Accept"] == "application/json"
assert (
downloader.session.headers["Authorization"]
== f"Bearer {bearertoken}"
)
bearertokenfile = join(downloaderfolder, "bearertoken.txt")
bearertoken = "12345"
with Download(bearer_token_file=bearertokenfile) as downloader:
assert downloader.session.headers["Accept"] == "application/json"
assert (
downloader.session.headers["Authorization"]
== f"Bearer {bearertoken}"
)
extraparamsyamltree = join(downloaderfolder, "extra_params_tree.yaml")
with Download(
extra_params_yaml=extraparamsyamltree, extra_params_lookup="mykey"
Expand All @@ -130,6 +145,19 @@ def test_init(self, monkeypatch, downloaderfolder):
monkeypatch.setenv("BASIC_AUTH", basicauth)
with Download() as downloader:
assert downloader.session.auth == ("user", "pass")
bearertoken = "98765"
monkeypatch.setenv("BEARER_TOKEN", bearertoken)
with pytest.raises(SessionError):
Download()
monkeypatch.delenv("BASIC_AUTH")
with Download() as downloader:
assert downloader.session.headers["Accept"] == "application/json"
assert (
downloader.session.headers["Authorization"]
== f"Bearer {bearertoken}"
)
monkeypatch.delenv("BEARER_TOKEN")
monkeypatch.setenv("BASIC_AUTH", basicauth)
with pytest.raises(SessionError):
Download(basic_auth="12345")
with pytest.raises(SessionError):
Expand Down

0 comments on commit 7649014

Please sign in to comment.