Skip to content

Commit

Permalink
Merge pull request #25 from PrivateAIM/23-integrate-into-hub-storage-api
Browse files Browse the repository at this point in the history
Add integration into Hub storage API
  • Loading branch information
mjugl authored Mar 7, 2024
2 parents 5b68cf6 + 659fd5c commit c1970cf
Show file tree
Hide file tree
Showing 17 changed files with 675 additions and 313 deletions.
14 changes: 6 additions & 8 deletions .env.example
Original file line number Diff line number Diff line change
@@ -1,17 +1,15 @@
# Local MinIO instance
# Hub endpoints
HUB__API_BASE_URL=https://api.privateaim.net
HUB__AUTH_BASE_URL=https://auth.privateaim.net
HUB__AUTH_USERNAME=foobar
HUB__AUTH_PASSWORD=sup3r_s3cr3t
# MinIO instance
MINIO__ENDPOINT=localhost:9000
MINIO__ACCESS_KEY=admin
MINIO__SECRET_KEY=s3cr3t_p4ssw0rd
MINIO__REGION=us-east-1
MINIO__BUCKET=flame
MINIO__USE_SSL=0
# Remote MinIO instance (should be different from local in productive setting)
REMOTE__ENDPOINT=localhost:9000
REMOTE__ACCESS_KEY=admin
REMOTE__SECRET_KEY=s3cr3t_p4ssw0rd
REMOTE__REGION=us-east-1
REMOTE__BUCKET=upload
REMOTE__USE_SSL=0
# OIDC config
OIDC__CERTS_URL=http://localhost:8080/realms/flame/protocol/openid-connect/certs
OIDC__CLIENT_ID_CLAIM_NAME=client_id
8 changes: 2 additions & 6 deletions .github/workflows/integration-test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -30,19 +30,15 @@ jobs:
run: poetry install
- name: Create local bucket
run: poetry run flame-result migrate --no-verify "$MINIO_ROOT_USER" "$MINIO_ROOT_PASSWORD" "$MINIO_LOCAL_BUCKET_NAME"
- name: Create remote bucket
run: poetry run flame-result migrate --no-verify "$MINIO_ROOT_USER" "$MINIO_ROOT_PASSWORD" "$MINIO_REMOTE_BUCKET_NAME"
- name: Run tests
run: poetry run pytest
env:
PYTEST__MINIO__ACCESS_KEY: ${{env.MINIO_ROOT_USER}}
PYTEST__MINIO__SECRET_KEY: ${{env.MINIO_ROOT_PASSWORD}}
PYTEST__MINIO__BUCKET: ${{env.MINIO_LOCAL_BUCKET_NAME}}
PYTEST__MINIO__ENDPOINT: localhost:9000
PYTEST__REMOTE__ACCESS_KEY: ${{env.MINIO_ROOT_USER}}
PYTEST__REMOTE__SECRET_KEY: ${{env.MINIO_ROOT_PASSWORD}}
PYTEST__REMOTE__BUCKET: ${{env.MINIO_REMOTE_BUCKET_NAME}}
PYTEST__REMOTE__ENDPOINT: localhost:9000
PYTEST__HUB__AUTH_USERNAME: ${{secrets.HUB_AUTH_USERNAME}}
PYTEST__HUB__AUTH_PASSWORD: ${{secrets.HUB_AUTH_PASSWORD}}

services:
minio:
Expand Down
6 changes: 2 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,8 @@ $ PYTEST__MINIO__ENDPOINT="localhost:9000" \
PYTEST__MINIO__ACCESS_KEY="admin" \
PYTEST__MINIO__SECRET_KEY="s3cr3t_p4ssw0rd" \
PYTEST__MINIO__BUCKET="flame" \
PYTEST__REMOTE__ENDPOINT="localhost:9000" \
PYTEST__REMOTE__ACCESS_KEY="admin" \
PYTEST__REMOTE__SECRET_KEY="s3cr3t_p4ssw0rd" \
PYTEST__REMOTE__BUCKET="upload" pytest
PYTEST__HUB__AUTH_USERNAME="XXXXXXXX" \
PYTEST__HUB__AUTH_PASSWORD="XXXXXXXX" pytest
```

OIDC does not need to be configured.
Expand Down
319 changes: 162 additions & 157 deletions poetry.lock

Large diffs are not rendered by default.

11 changes: 10 additions & 1 deletion project/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,18 @@ class OIDCConfig(BaseModel):
model_config = ConfigDict(frozen=True)


class HubConfig(BaseModel):
api_base_url: HttpUrl = "https://api.privateaim.net"
auth_base_url: HttpUrl = "https://auth.privateaim.net"
auth_username: str
auth_password: str

model_config = ConfigDict(frozen=True)


class Settings(BaseSettings):
hub: HubConfig
minio: MinioBucketConfig
remote: MinioBucketConfig
oidc: OIDCConfig

model_config = SettingsConfigDict(
Expand Down
43 changes: 37 additions & 6 deletions project/dependencies.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import json
import logging
import time
from functools import lru_cache
from typing import Annotated

Expand All @@ -12,10 +13,15 @@
from starlette import status

from project.config import Settings, MinioBucketConfig
from project.hub import AccessToken, AuthWrapper

security = HTTPBearer()
logger = logging.getLogger(__name__)

# TODO this doesn't consider the fact that an access token may be invalidated for any reason other than expiration
_access_token: AccessToken | None = None
_access_token_retrieved_at: int


@lru_cache
def get_settings():
Expand Down Expand Up @@ -56,12 +62,6 @@ def get_local_minio(
return __create_minio_from_config(settings.minio)


def get_remote_minio(
settings: Annotated[Settings, Depends(get_settings)],
):
return __create_minio_from_config(settings.remote)


def get_client_id(
settings: Annotated[Settings, Depends(get_settings)],
jwks: Annotated[jwk.JWKSet, Depends(get_auth_jwks)],
Expand All @@ -87,3 +87,34 @@ def get_client_id(
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail="JWT is malformed"
)


def __obtain_new_access_token(auth: AuthWrapper, username: str, password: str):
global _access_token, _access_token_retrieved_at

_access_token = auth.acquire_access_token_with_password(username, password)
_access_token_retrieved_at = int(time.time())


def get_auth_wrapper(settings: Annotated[Settings, Depends(get_settings)]):
return AuthWrapper(str(settings.hub.auth_base_url))


def get_access_token(
settings: Annotated[Settings, Depends(get_settings)],
auth_wrapper: Annotated[AuthWrapper, Depends(get_auth_wrapper)],
) -> AccessToken:
global _access_token, _access_token_retrieved_at

if _access_token is None:
__obtain_new_access_token(
auth_wrapper, settings.hub.auth_username, settings.hub.auth_password
)

# TODO configurable leeway?
if int(time.time()) < _access_token_retrieved_at - 3600:
__obtain_new_access_token(
auth_wrapper, settings.hub.auth_username, settings.hub.auth_password
)

return _access_token
221 changes: 221 additions & 0 deletions project/hub.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,221 @@
from io import BytesIO
from typing import NamedTuple
from urllib.parse import urljoin

import httpx
from starlette import status


class AccessToken(NamedTuple):
access_token: str
expires_in: int
token_type: str
scope: str
refresh_token: str


class Project(NamedTuple):
id: str
name: str


class Analysis(NamedTuple):
id: str
name: str


class BucketFile(NamedTuple):
id: str
name: str
bucket_id: str


class AnalysisFile(NamedTuple):
id: str
name: str
type: str
bucket_file_id: str


class Bucket(NamedTuple):
id: str
name: str


class AuthWrapper:
"""Simple wrapper around the password-based token grant of the central AuthUp instance."""

def __init__(self, base_url: str):
self.base_url = base_url

def acquire_access_token_with_password(
self, username: str, password: str
) -> AccessToken:
"""Acquire an access token using the given username and password."""
r = httpx.post(
urljoin(self.base_url, "/token"),
json={
"grant_type": "password",
"username": username,
"password": password,
},
).raise_for_status()
j = r.json()

return AccessToken(
access_token=j["access_token"],
expires_in=j["expires_in"],
token_type=j["token_type"],
scope=j["scope"],
refresh_token=j["refresh_token"],
)


class ApiWrapper:
"""Simple wrapper around the central hub API. The wrapper does NOT check the access token validity."""

def __init__(self, base_url: str, access_token: str):
self.base_url = base_url
self.access_token = access_token

def __auth_header(self):
return {"Authorization": f"Bearer {self.access_token}"}

def create_project(self, name: str) -> Project:
"""Create a project with the given name."""
r = httpx.post(
urljoin(self.base_url, "/projects"),
headers=self.__auth_header(),
json={"name": name},
).raise_for_status()
j = r.json()

return Project(
id=j["id"],
name=j["name"],
)

def create_analysis(self, name: str, project_id: str) -> Analysis:
"""Create an analysis with the given name and assign it to the given project."""
r = httpx.post(
urljoin(self.base_url, "/analyses"),
headers=self.__auth_header(),
json={
"name": name,
"project_id": project_id,
},
).raise_for_status()
j = r.json()

return Analysis(
id=j["id"],
name=j["name"],
)

def get_bucket(self, bucket_name: str) -> Bucket | None:
"""Get the bucket associated with the given name."""
r = httpx.get(
urljoin(self.base_url, f"/storage/buckets/{bucket_name}"),
headers=self.__auth_header(),
)

if r.status_code == status.HTTP_404_NOT_FOUND:
return None

# catch any other unexpected status
r.raise_for_status()
j = r.json()

return Bucket(
id=j["id"],
name=j["name"],
)

def get_bucket_file(self, bucket_file_id: str) -> BucketFile | None:
"""Get the file associated with the given bucket file ID."""
r = httpx.get(
urljoin(self.base_url, f"/storage/bucket-files/{bucket_file_id}"),
headers=self.__auth_header(),
).raise_for_status()

if r.status_code == status.HTTP_404_NOT_FOUND:
return None

r.raise_for_status()
j = r.json()

return BucketFile(
id=j["id"],
name=j["name"],
bucket_id=j["bucket_id"],
)

def upload_to_bucket(
self,
bucket_name: str,
file_name: str,
file: BytesIO,
content_type: str = "application/octet-stream",
) -> list[BucketFile]:
"""Upload a file to the bucket associated with the given name. Content type is optional and is set
to application/octet-stream by default."""
r = httpx.post(
urljoin(self.base_url, f"/storage/buckets/{bucket_name}/upload"),
headers=self.__auth_header(),
files={
"file": (file_name, file, content_type),
},
).raise_for_status()
j = r.json()

return [
BucketFile(
id=b["id"],
name=b["name"],
bucket_id=b["bucket_id"],
)
for b in j["data"]
]

def link_file_to_analysis(
self, analysis_id: str, bucket_file_id: str, bucket_file_name: str
) -> AnalysisFile:
"""Link the file associated with the given ID and name to the analysis associated with the given ID.
Currently, this function only supports linking result files."""
r = httpx.post(
urljoin(self.base_url, "/analysis-files"),
headers=self.__auth_header(),
json={
"analysis_id": analysis_id,
"type": "RESULT",
"bucket_file_id": bucket_file_id,
"name": bucket_file_name,
"root": True,
},
).raise_for_status()
j = r.json()

return AnalysisFile(
id=j["id"],
name=j["name"],
type=j["type"],
bucket_file_id=j["bucket_file_id"],
)

def get_analysis_files(self) -> list[AnalysisFile]:
"""List all analysis files."""
r = httpx.get(
urljoin(self.base_url, "/analysis-files"),
headers=self.__auth_header(),
).raise_for_status()
j = r.json()

return [
AnalysisFile(
id=f["id"],
name=f["name"],
type=f["type"],
bucket_file_id=f["bucket_file_id"],
)
for f in j["data"]
]
Loading

0 comments on commit c1970cf

Please sign in to comment.