feat: ✨ add support for reading client credentials from the local H2O…

… CLI config file (#68)

* feat: ✨ add support for reading clientcredentials form the local H2O CLI config

 - implment config path lookusp
 - implement config loading
 - extend credentials loading to accept tokens loaded from the config
 - extend main discover functions to accept config path and utilize

RESOLVES h2oai/cloud-discovery#437

* REVIEW: tomli python_version squote

pyproject.toml

@@ -14,7 +14,10 @@ classifiers = [
   "Topic :: Security",
   "Topic :: Utilities",
 ]
-dependencies = ["httpx>=0.16"]
+dependencies = [
+  "httpx>=0.16",
+  "tomli >= 1.1.0 ; python_version < '3.11'",
+]
 description = 'H2O Cloud Discovery Python CLient'
 keywords = []
 license = "Apache-2.0"

src/h2o_discovery/__init__.py

@@ -1,8 +1,11 @@
 import dataclasses
+import os
 from typing import Optional
+from typing import Union
 
 from h2o_discovery import model
 from h2o_discovery._internal import client
+from h2o_discovery._internal import config
 from h2o_discovery._internal import load
 from h2o_discovery._internal import lookup
 from h2o_discovery._version import __version__  # noqa: F401
@@ -11,44 +14,79 @@
 
 
 def discover(
-    environment: Optional[str] = None, discovery_address: Optional[str] = None
+    environment: Optional[str] = None,
+    discovery_address: Optional[str] = None,
+    config_path: Optional[Union[str, bytes, os.PathLike]] = None,
 ) -> Discovery:
     """Obtains and returns a Discovery object from the discovery service.
 
-    Both arguments are optional. If neither is provided, the environment variable
-    H2O_CLOUD_ENVIRONMENT is used. If that is not set, the environment variable
-    H2O_CLOUD_DISCOVERY is used. If that is not set, a LookupError is raised.
+    All arguments are optional. Discovery determined with the following precedence:
+        - discovery_address parameter
+        - H2O_CLOUD_DISCOVERY environment variable
+        - environment parameter
+        - H2O_CLOUD_ENVIRONMENT environment variable
+
+    Config path is determined with the following precedence:
+        - config_path parameter
+        - H2OCONFIG environment variable
+        - default H2O CLI configuration configuration path
+        - if default path does not exist, no config is loaded
 
     Args:
         environment: The H2O Cloud environment URL to use (e.g. https://cloud.h2o.ai).
         discovery_address: The address of the discovery service.
+        config_path: The path to the H2O CLI configuration file. Only used for
+            credentials lookup.
 
+    Raises:
+        LookupError: If the URI cannot be determined.
     """
     uri = lookup.determine_uri(environment, discovery_address)
+    cfg = config.load_config(lookup.determine_local_config_path(config_path))
 
     discovery = load.load_discovery(client.Client(uri))
-    credentials = load.load_credentials(clients=discovery.clients)
+    credentials = load.load_credentials(
+        clients=discovery.clients, config_tokens=cfg.tokens
+    )
 
     return dataclasses.replace(discovery, credentials=credentials)
 
 
 async def discover_async(
-    environment: Optional[str] = None, discovery_address: Optional[str] = None
+    environment: Optional[str] = None,
+    discovery_address: Optional[str] = None,
+    config_path: Optional[Union[str, bytes, os.PathLike]] = None,
 ) -> Discovery:
     """Obtains and returns a Discovery object from the discovery service.
 
-    Both arguments are optional. If neither is provided, the environment variable
-    H2O_CLOUD_ENVIRONMENT is used. If that is not set, the environment variable
-    H2O_CLOUD_DISCOVERY is used. If that is not set, a LookupError is raised.
+    All arguments are optional. Discovery determined with the following precedence:
+        - discovery_address parameter
+        - H2O_CLOUD_DISCOVERY environment variable
+        - environment parameter
+        - H2O_CLOUD_ENVIRONMENT environment variable
+
+    Config path is determined with the following precedence:
+        - config_path parameter
+        - H2OCONFIG environment variable
+        - default H2O CLI configuration configuration path
+        - if default path does not exist, no config is loaded
 
     Args:
         environment: The H2O Cloud environment URL to use (e.g. https://cloud.h2o.ai).
         discovery_address: The address of the discovery service.
+        config_path: The path to the H2O CLI configuration file. Only used for
+            credentials lookup.
 
+    Raises:
+        LookupError: If the URI cannot be determined.
     """
+
     uri = lookup.determine_uri(environment, discovery_address)
+    cfg = config.load_config(lookup.determine_local_config_path(config_path))
 
     discovery = await load.load_discovery_async(client.AsyncClient(uri))
-    credentials = load.load_credentials(clients=discovery.clients)
+    credentials = load.load_credentials(
+        clients=discovery.clients, config_tokens=cfg.tokens
+    )
 
     return dataclasses.replace(discovery, credentials=credentials)

src/h2o_discovery/_internal/compat.py

@@ -0,0 +1,11 @@
+"""Helper module for importing backported packages. So that the rest of the
+modules is not cluttered with conditional imports.
+"""
+
+import sys
+
+# We use version check instead of try/except because mypy does not like it.
+if sys.version_info >= (3, 11):
+    import tomllib  # noqa: F401
+else:
+    import tomli as tomllib  # noqa: F401

src/h2o_discovery/_internal/config.py

@@ -0,0 +1,45 @@
+import dataclasses
+import types
+from typing import Mapping
+from typing import Optional
+
+from h2o_discovery._internal.compat import tomllib
+
+
+def _empty_tokens_factory() -> Mapping[str, str]:
+    return types.MappingProxyType({})
+
+
+@dataclasses.dataclass(frozen=True)
+class Config:
+    """Internal representation of the H2O CLI Configuration."""
+
+    #: Map of found tokens in the configuration file. in the `{"client-id": "token"}`
+    #: format.
+    tokens: Mapping[str, str] = dataclasses.field(default_factory=_empty_tokens_factory)
+
+
+def load_config(path: Optional[str] = None) -> Config:
+    """Loads the H2O CLI config from the specified path.
+
+    If no path is specified, an empty configuration is returned.
+    """
+
+    if not path:
+        return Config()
+
+    with open(path, "rb") as f:
+        data = tomllib.load(f)
+
+    client_id = data.get("ClientID")
+    token = data.get("Token")
+    platform_client_id = data.get("PlatformClientID")
+    platform_token = data.get("PlatformToken")
+
+    tokens = {}
+    if client_id is not None and token is not None:
+        tokens[client_id] = token
+    if platform_client_id is not None and platform_token is not None:
+        tokens[platform_client_id] = platform_token
+
+    return Config(tokens=types.MappingProxyType(tokens))

src/h2o_discovery/_internal/load.py

@@ -2,6 +2,7 @@
 import types
 from typing import Iterable
 from typing import Mapping
+from typing import Optional
 
 from h2o_discovery import model
 from h2o_discovery._internal import client
@@ -26,17 +27,21 @@ async def load_discovery_async(cl: client.AsyncClient) -> model.Discovery:
 
 
 def load_credentials(
-    clients: Mapping[str, model.Client]
+    clients: Mapping[str, model.Client], config_tokens: Optional[Mapping[str, str]]
 ) -> Mapping[str, model.Credentials]:
-    """Loads client credentials from the environment."""
-    out = {}
+    """Loads client credentials from the environment or tokens loaded from the
+    config.
+    """
+    tokens: Mapping[str, str] = {}
+    if config_tokens is not None:
+        tokens = config_tokens
 
-    for name in clients.keys():
+    out = {}
+    for name, cl in clients.items():
         env_name = f"H2O_CLOUD_CLIENT_{name.upper()}_TOKEN"
-        token = os.environ.get(env_name)
+        token = os.environ.get(env_name) or tokens.get(cl.oauth2_client_id)
         if token:
             out[name] = model.Credentials(client=name, refresh_token=token)
-
     return types.MappingProxyType(out)
 
 

src/h2o_discovery/_internal/lookup.py

@@ -1,8 +1,10 @@
 import os
 from typing import Optional
+from typing import Union
 import urllib.parse
 
 _WELL_KNOWN_PATH = ".ai.h2o.cloud.discovery"
+_DEFAULT_LOCAL_CONFIG_PATH = "~/.h2oai/h2o-cli-config.toml"
 
 
 def determine_uri(
@@ -37,3 +39,22 @@ def determine_uri(
 
 def _discovery_uri_from_environment(environment: str):
     return urllib.parse.urljoin(environment + "/", _WELL_KNOWN_PATH)
+
+
+def determine_local_config_path(
+    config_path: Optional[Optional[Union[str, bytes, os.PathLike]]] = None
+) -> Optional[str]:
+    """Uses passed parameter, environment variable and H2O CLI default to get the
+    path to the local config file.
+    """
+    if config_path is not None:
+        return str(os.fspath(config_path))
+
+    config_path = os.environ.get("H2OCONFIG")
+    if config_path is not None:
+        return config_path
+
+    local_config_path = os.path.expanduser(_DEFAULT_LOCAL_CONFIG_PATH)
+    if not os.path.isfile(local_config_path):
+        return None
+    return local_config_path

tests/_internal/load/test_load_credentials.py

@@ -10,17 +10,27 @@ def test_load_credentials(monkeypatch):
             display_name="Env Client",
             oauth2_client_id="env-client-id",
         ),
+        "config_client": model.Client(
+            name="clients/config_client",
+            display_name="Config Client",
+            oauth2_client_id="config-client-id",
+        ),
         "extra_client": model.Client(
             name="clients/extra_client",
             display_name="Extra Client",
             oauth2_client_id="extra-client-id",
         ),
     }
+    tokens = {"config-client-id": "config-token"}
     monkeypatch.setenv("H2O_CLOUD_CLIENT_ENV_CLIENT_TOKEN", "env-token")
 
     # When
-    result = load.load_credentials(clients=clients)
+    result = load.load_credentials(clients=clients, config_tokens=tokens)
+
     # Then
     assert result == {
-        "env_client": model.Credentials(client="env_client", refresh_token="env-token")
+        "env_client": model.Credentials(client="env_client", refresh_token="env-token"),
+        "config_client": model.Credentials(
+            client="config_client", refresh_token="config-token"
+        ),
     }

tests/_internal/lookup/test_determine_local_config_path.py

@@ -0,0 +1,55 @@
+from h2o_discovery._internal import lookup
+
+
+def test_determine_local_config_path_param():
+    # When
+    path = lookup.determine_local_config_path("test-path")
+
+    # Then
+    assert path == "test-path"
+
+
+def test_determine_local_config_path_env(monkeypatch):
+    # Given
+    monkeypatch.setenv("H2OCONFIG", "test-path-from-env")
+
+    # When
+    path = lookup.determine_local_config_path()
+
+    # Then
+    assert path == "test-path-from-env"
+
+
+def test_determine_local_config_path_default_location_does_not_exist(
+    monkeypatch, tmp_path
+):
+    # Given
+    home = tmp_path / "home" / "test-user"
+    home.mkdir(parents=True)
+    monkeypatch.setenv("HOME", str(home))
+
+    # When
+    path = lookup.determine_local_config_path()
+
+    # Then
+    assert path is None
+
+
+def test_determine_local_config_path_default_location_does_exists(
+    monkeypatch, tmp_path
+):
+    # Given
+    home = tmp_path / "home" / "test-user"
+    home.mkdir(parents=True)
+    monkeypatch.setenv("HOME", str(home))
+
+    config_dir = home / ".h2oai"
+    config_dir.mkdir(parents=True)
+    config_file = config_dir / "h2o-cli-config.toml"
+    config_file.touch()
+
+    # When
+    path = lookup.determine_local_config_path()
+
+    # Then
+    assert path == str(config_file)

tests/_internal/test_config.py

@@ -0,0 +1,56 @@
+import pytest
+
+from h2o_discovery._internal import config
+
+
+def config_test_cases():
+    yield pytest.param("", config.Config(tokens={}), id="empty config")
+    yield pytest.param(
+        """
+        ClientID = "client-id"
+        Token = "TestToken"
+        PlatformClientID = "platform-client-id"
+        PlatformToken = "TestPlatformToken"
+        """,
+        config.Config(
+            tokens={"client-id": "TestToken", "platform-client-id": "TestPlatformToken"}
+        ),
+        id="full config",
+    )
+    yield pytest.param(
+        """
+        ClientID = "client-id"
+        Token = "TestToken"
+        """,
+        config.Config(tokens={"client-id": "TestToken"}),
+        id="cli token only",
+    )
+    yield pytest.param(
+        """
+        PlatformClientID = "platform-client-id"
+        PlatformToken = "TestPlatformToken"
+        """,
+        config.Config(tokens={"platform-client-id": "TestPlatformToken"}),
+        id="platform token only",
+    )
+
+
+@pytest.mark.parametrize("config_content,expected_results", config_test_cases())
+def test_load_config(config_content, expected_results, tmp_path):
+    # Given
+    config_file = tmp_path / "config.toml"
+    config_file.write_text(config_content)
+
+    # When
+    result = config.load_config(config_file)
+
+    # Then
+    assert result == expected_results
+
+
+def test_load_config_no_path(tmp_path):
+    # When
+    result = config.load_config()
+
+    # Then
+    assert result == config.Config()