[QI2-1211] Move Qiskit Plugin's API subpackage to compute-api-client

Author: eliasKA

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "qi-compute-api-client"
-version = "0.42.0"
+version = "0.43.0"
 description = "An API client for the Compute Job Manager of Quantum Inspire."
 license = "Apache-2.0"
 authors = ["Quantum Inspire <[email protected]>"]
@@ -18,15 +18,18 @@ classifiers = [
 ]
 packages = [
     { include = "compute_api_client" },
+    { include = "qi2_shared" },
 ]
-exclude = ["compute_api_client/test"]
+exclude = ["compute_api_client/test", "tests"]
 
 [tool.poetry.dependencies]
 # Should not exceed package python version, so be conservative
 aiohttp = "^3.10.5"
+pydantic = "^2.10.4"
 python = "^3.8"
 python-dateutil = "^2.8.2"
 urllib3 = "^2.0.0"
+requests = "^2.32.3"
 
 [tool.poetry.group.dev.dependencies]
 pytest = {extras = ["toml"], version = "^8.0.0"}

pyproject.toml

@@ -0,0 +1,65 @@
+import time
+from typing import Any, Tuple, cast
+
+import requests
+
+from qi2_shared.settings import ApiSettings, TokenInfo, Url
+
+
+class AuthorisationError(Exception):
+    """Indicates that the authorisation permanently went wrong."""
+    pass
+
+
+class IdentityProvider:
+    """Class for interfacing with the IdentityProvider."""
+
+    def __init__(self, well_known_endpoint: str):
+        self._well_known_endpoint = well_known_endpoint
+        self._token_endpoint, self._device_endpoint = self._get_endpoints()
+        self._headers = {"Content-Type": "application/x-www-form-urlencoded"}
+
+    def _get_endpoints(self) -> Tuple[str, str]:
+        response = requests.get(self._well_known_endpoint)
+        response.raise_for_status()
+        config = response.json()
+        return config["token_endpoint"], config["device_authorization_endpoint"]
+
+    def refresh_access_token(self, client_id: str, refresh_token: str) -> dict[str, Any]:
+        data = {
+            "grant_type": "refresh_token",
+            "client_id": client_id,
+            "refresh_token": refresh_token,
+        }
+        response = requests.post(self._token_endpoint, headers=self._headers, data=data)
+        response.raise_for_status()
+        return cast(dict[str, Any], response.json())
+
+
+class OauthDeviceSession:
+    """Class for storing OAuth session information and refreshing tokens when needed."""
+
+    def __init__(self, host: Url, settings: ApiSettings, identity_provider: IdentityProvider):
+        self._api_settings = settings
+        _auth_settings = settings.auths[host]
+        self._host = host
+        self._client_id = _auth_settings.client_id
+        self._token_info = _auth_settings.tokens
+        self._refresh_time_reduction = 5  # the number of seconds to refresh the expiration time
+        self._identity_provider = identity_provider
+
+    def refresh(self) -> TokenInfo:
+        if self._token_info is None:
+            raise AuthorisationError("You should authenticate first before you can refresh")
+
+        if self._token_info.access_expires_at > time.time() + self._refresh_time_reduction:
+            return self._token_info
+
+        try:
+            self._token_info = TokenInfo(
+                **self._identity_provider.refresh_access_token(self._client_id, self._token_info.refresh_token)
+            )
+            self._api_settings.store_tokens(self._host, self._token_info)
+            return self._token_info
+        except requests.HTTPError as e:
+            raise AuthorisationError(f"An error occurred during token refresh: {e}")

qi2_shared/__init__.py

@@ -0,0 +1,52 @@
+from typing import Any, Optional
+
+import compute_api_client
+
+from qi2_shared.authentication import IdentityProvider, OauthDeviceSession
+from qi2_shared.settings import ApiSettings
+
+
+class Configuration(compute_api_client.Configuration):  # type: ignore[misc]
+    """Original Configuration class in compute_api_client does not handle refreshing bearer tokens, so we need to add
+    some functionality."""
+
+    def __init__(self, host: str, oauth_session: OauthDeviceSession, **kwargs: Any):
+        self._oauth_session = oauth_session
+        super().__init__(host=host, **kwargs)
+
+    def auth_settings(self) -> Any:
+        token_info = self._oauth_session.refresh()
+        self.access_token = token_info.access_token
+        return super().auth_settings()
+
+
+_config: Optional[Configuration] = None
+
+
+def connect() -> None:
+    """Set connection configuration for the Quantum Inspire API.
+
+    Call after logging in with the CLI. Will remove old configuration.
+    """
+    global _config
+    settings = ApiSettings.from_config_file()
+
+    tokens = settings.auths[settings.default_host].tokens
+
+    if tokens is None:
+        raise ValueError("No access token found for the default host. Please connect to Quantum Inspire using the CLI.")
+
+    host = settings.default_host
+    _config = Configuration(
+        host=host,
+        oauth_session=OauthDeviceSession(host, settings, IdentityProvider(settings.auths[host].well_known_endpoint)),
+    )
+
+
+def config() -> Configuration:
+    global _config
+    if _config is None:
+        connect()
+
+    assert _config is not None
+    return _config

qi2_shared/authentication.py

@@ -0,0 +1,44 @@
+from typing import Any, Awaitable, Callable, Generic, List, Optional, TypeVar, Union, cast
+
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+PageType = TypeVar("PageType")
+ItemType = TypeVar("ItemType")
+
+
+class PageInterface(BaseModel, Generic[ItemType]):
+    """The page models in the generated API client don't inherit from a common base class, so we have to trick the
+    typing system a bit with this fake base class."""
+
+    items: List[ItemType]
+    total: Optional[Annotated[int, Field(strict=True, ge=0)]]
+    page: Optional[Annotated[int, Field(strict=True, ge=1)]]
+    size: Optional[Annotated[int, Field(strict=True, ge=1)]]
+    pages: Optional[Annotated[int, Field(strict=True, ge=0)]] = None
+
+
+class PageReader(Generic[PageType, ItemType]):
+    """Helper class for reading fastapi-pagination style pages returned by the compute_api_client."""
+
+    async def get_all(self, api_call: Callable[..., Awaitable[PageType]], **kwargs: Any) -> List[ItemType]:
+        """Get all items from an API call that supports paging."""
+        items: List[ItemType] = []
+        page = 1
+
+        while True:
+            response = cast(PageInterface[ItemType], await api_call(page=page, **kwargs))
+
+            items.extend(response.items)
+            page += 1
+            if response.pages is None or page > response.pages:
+                break
+        return items
+
+    async def get_single(self, api_call: Callable[..., Awaitable[PageType]], **kwargs: Any) -> Union[ItemType, None]:
+        """Get a single item from an API call that supports paging."""
+        response = cast(PageInterface[ItemType], await api_call(**kwargs))
+        if len(response.items) > 1:
+            raise RuntimeError(f"Response contains more than one item -> {kwargs}.")
+
+        return response.items[0] if response.items else None

qi2_shared/client.py

@@ -0,0 +1,68 @@
+"""Module containing the handler for the Quantum Inspire persistent configuration."""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import Dict, Optional
+
+from pydantic import BaseModel, BeforeValidator, Field, HttpUrl
+from typing_extensions import Annotated
+
+Url = Annotated[str, BeforeValidator(lambda value: str(HttpUrl(value)).rstrip("/"))]
+API_SETTINGS_FILE = Path.joinpath(Path.home(), ".quantuminspire", "config.json")
+
+
+class TokenInfo(BaseModel):
+    """A pydantic model for storing all information regarding oauth access and refresh tokens."""
+
+    access_token: str
+    expires_in: int  # [s]
+    refresh_token: str
+    refresh_expires_in: Optional[int] = None  # [s]
+    generated_at: float = Field(default_factory=time.time)
+
+    @property
+    def access_expires_at(self) -> float:
+        """Unix timestamp containing the time when the access token will expire."""
+        return self.generated_at + self.expires_in
+
+
+class AuthSettings(BaseModel):
+    """Pydantic model for storing all auth related settings for a given host."""
+
+    client_id: str
+    code_challenge_method: str
+    code_verifyer_length: int
+    well_known_endpoint: Url
+    tokens: Optional[TokenInfo]
+    team_member_id: Optional[int]
+
+
+class ApiSettings(BaseModel):
+    """The settings class for the Quantum Inspire persistent configuration."""
+
+    auths: Dict[Url, AuthSettings]
+    default_host: Url
+
+    def store_tokens(self, host: Url, tokens: TokenInfo, path: Path = API_SETTINGS_FILE) -> None:
+        """Stores the team_member_id, access and refresh tokens in the config.json file.
+
+        Args:
+            host: The hostname of the API for which the tokens are intended.
+            tokens: OAuth access and refresh tokens.
+            path: The path to the config.json file. Defaults to API_SETTINGS_FILE.
+        Returns:
+            None
+        """
+        self.auths[host].tokens = tokens
+        path.write_text(self.model_dump_json(indent=2))
+
+    @classmethod
+    def from_config_file(cls, path: Path = API_SETTINGS_FILE) -> ApiSettings:
+        """Load the configuration from a file."""
+        if not path.is_file():
+            raise FileNotFoundError("No configuration file found. Please connect to Quantum Inspire using the CLI.")
+
+        api_settings = path.read_text()
+        return ApiSettings.model_validate_json(api_settings)

qi2_shared/pagination.py

@@ -0,0 +1,26 @@
+import pytest
+
+from qi2_shared.settings import AuthSettings, TokenInfo
+
+
+@pytest.fixture
+def token_info() -> TokenInfo:
+    return TokenInfo(
+        access_token="access_token",
+        expires_in=100,
+        refresh_token="refresh_token",
+        refresh_expires_in=1000,
+        generated_at=1,
+    )
+
+
+@pytest.fixture
+def auth_settings(token_info: TokenInfo) -> AuthSettings:
+    return AuthSettings(
+        client_id="client_id",
+        code_challenge_method="code_challenge_method",
+        code_verifyer_length=1,
+        well_known_endpoint="https://host.com/well-known-endpoint",
+        tokens=token_info,
+        team_member_id=1,
+    )

qi2_shared/settings.py

@@ -0,0 +1,108 @@
+import time
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+import responses
+from responses import matchers
+
+from qi2_shared.authentication import AuthorisationError, IdentityProvider, OauthDeviceSession
+from qi2_shared.settings import ApiSettings, AuthSettings, TokenInfo
+
+
+@pytest.fixture
+def identity_provider_mock() -> MagicMock:
+    return MagicMock(spec=IdentityProvider)
+
+
+@pytest.fixture
+def api_settings_mock(auth_settings: AuthSettings) -> MagicMock:
+    api_settings = MagicMock(spec=ApiSettings)
+    api_settings.default_host = "https://host.com"
+    api_settings.auths = {api_settings.default_host: auth_settings}
+    return api_settings
+
+
+def test_oauth_device_session_refresh_no_token(api_settings_mock: MagicMock, identity_provider_mock: MagicMock) -> None:
+    # Arrange
+    api_settings_mock.auths[api_settings_mock.default_host].tokens = None
+    session = OauthDeviceSession("https://host.com", api_settings_mock, identity_provider_mock)
+
+    # Act & Assert
+    with pytest.raises(AuthorisationError):
+        session.refresh()
+
+
+def test_oauth_device_session_refresh_token_not_expired(
+    api_settings_mock: MagicMock, identity_provider_mock: MagicMock
+) -> None:
+    # Arrange
+    auth_settings = api_settings_mock.auths[api_settings_mock.default_host]
+    auth_settings.tokens.generated_at = time.time()
+    session = OauthDeviceSession("https://host.com", api_settings_mock, identity_provider_mock)
+
+    # Act
+    token_info = session.refresh()
+
+    # Assert
+    assert token_info == auth_settings.tokens
+
+    identity_provider_mock.refresh_access_token.assert_not_called()
+
+
+def test_oauth_device_session_refresh_token_expired(
+    api_settings_mock: MagicMock, identity_provider_mock: MagicMock
+) -> None:
+    # Arrange
+    session = OauthDeviceSession("https://host.com", api_settings_mock, identity_provider_mock)
+    new_token_info: dict[str, Any] = {
+        "access_token": "new_access_token",
+        "expires_in": 100,
+        "refresh_token": "new_refresh_token",
+        "refresh_expires_in": 1000,
+        "generated_at": time.time(),
+    }
+
+    identity_provider_mock.refresh_access_token.return_value = new_token_info
+
+    # Act
+    token_info = session.refresh()
+
+    # Assert
+    assert token_info == TokenInfo(**new_token_info)
+
+    identity_provider_mock.refresh_access_token.assert_called_once_with("client_id", "refresh_token")
+    api_settings_mock.store_tokens.assert_called_once_with("https://host.com", token_info)
+
+
+@responses.activate
+def test_identity_provider_refresh_access_token() -> None:
+    # Arrange
+    token_info = {"token": "something", "some": "other_data"}
+    client_id = "some_client"
+    old_refresh_token = "old_token"
+
+    responses.get(
+        "https://host.com/well-known-endpoint",
+        json={
+            "token_endpoint": "https://host.com/token-endpoint",
+            "device_authorization_endpoint": "https://host.com/device-endpoint",
+        },
+    )
+    responses.post(
+        "https://host.com/token-endpoint",
+        json=token_info,
+        match=[
+            matchers.urlencoded_params_matcher(
+                {"grant_type": "refresh_token", "client_id": client_id, "refresh_token": old_refresh_token}
+            )
+        ],
+    )
+
+    # Act
+    provider = IdentityProvider("https://host.com/well-known-endpoint")
+
+    token = provider.refresh_access_token(client_id, old_refresh_token)
+
+    # Assert
+    assert token == token_info