diff --git a/.github/.backend_git_ref b/.github/.backend_git_ref index 88d050b..ba2906d 100644 --- a/.github/.backend_git_ref +++ b/.github/.backend_git_ref @@ -1 +1 @@ -main \ No newline at end of file +main diff --git a/README.md b/README.md index a4dd640..ae71a9e 100644 --- a/README.md +++ b/README.md @@ -137,7 +137,7 @@ python3 -m mypy tests Using the config file `mypy.ini`, you can suppress missing stub errors for external libraries. You can ignore a library by adding two lines to the config file. For example, suppressing matplotlib would look like this: -``` +```ini [mypy-matplotlib.*] ignore_missing_imports = True diff --git a/geoengine/__init__.py b/geoengine/__init__.py index eddef6c..296cc07 100644 --- a/geoengine/__init__.py +++ b/geoengine/__init__.py @@ -19,7 +19,7 @@ from .layers import Layer, LayerCollection, LayerListing, LayerCollectionListing, \ LayerId, LayerCollectionId, LayerProviderId, \ layer_collection, layer -from .ml import register_ml_model, MlModelConfig +from .ml import register_ml_model, MlModelConfig, MlModelName from .permissions import add_permission, remove_permission, add_role, remove_role, assign_role, revoke_role, \ ADMIN_ROLE_ID, REGISTERED_USER_ROLE_ID, ANONYMOUS_USER_ROLE_ID, Permission, Resource, UserId, RoleId from .tasks import Task, TaskId diff --git a/geoengine/datasets.py b/geoengine/datasets.py index 3a9d953..036cd8d 100644 --- a/geoengine/datasets.py +++ b/geoengine/datasets.py @@ -257,7 +257,7 @@ def to_api_enum(self) -> geoengine_openapi_client.OgrSourceErrorSpec: class DatasetName: - '''A wrapper for a dataset id''' + '''A wrapper for a dataset name''' __dataset_name: str @@ -266,7 +266,7 @@ def __init__(self, dataset_name: str) -> None: @classmethod def from_response(cls, response: geoengine_openapi_client.CreateDatasetHandler200Response) -> DatasetName: - '''Parse a http response to an `DatasetId`''' + '''Parse a http response to an `DatasetName`''' return DatasetName(response.dataset_name) def __str__(self) -> str: @@ -276,7 +276,7 @@ def __repr__(self) -> str: return str(self) def __eq__(self, other) -> bool: - '''Checks if two dataset ids are equal''' + '''Checks if two dataset names are equal''' if not isinstance(other, self.__class__): return False diff --git a/geoengine/ml.py b/geoengine/ml.py index de35530..adcdeda 100644 --- a/geoengine/ml.py +++ b/geoengine/ml.py @@ -2,9 +2,11 @@ Util functions for machine learning ''' +from __future__ import annotations from pathlib import Path import tempfile from dataclasses import dataclass +import geoengine_openapi_client.models from onnx import TypeProto, TensorProto, ModelProto from onnx.helper import tensor_dtype_to_string from geoengine_openapi_client.models import MlModelMetadata, MlModel, RasterDataType @@ -23,10 +25,42 @@ class MlModelConfig: description: str = "My Ml Model Description" +class MlModelName: + '''A wrapper for an MlModel name''' + + __ml_model_name: str + + def __init__(self, ml_model_name: str) -> None: + self.__ml_model_name = ml_model_name + + @classmethod + def from_response(cls, response: geoengine_openapi_client.models.MlModelNameResponse) -> MlModelName: + '''Parse a http response to an `DatasetName`''' + return MlModelName(response.ml_model_name) + + def __str__(self) -> str: + return self.__ml_model_name + + def __repr__(self) -> str: + return str(self) + + def __eq__(self, other) -> bool: + '''Checks if two dataset names are equal''' + if not isinstance(other, self.__class__): + return False + + return self.__ml_model_name == other.__ml_model_name # pylint: disable=protected-access + + def to_api_dict(self) -> geoengine_openapi_client.models.MlModelNameResponse: + return geoengine_openapi_client.models.MlModelNameResponse( + ml_model_name=str(self.__ml_model_name) + ) + + def register_ml_model(onnx_model: ModelProto, model_config: MlModelConfig, upload_timeout: int = 3600, - register_timeout: int = 60): + register_timeout: int = 60) -> MlModelName: '''Uploads an onnx file and registers it as an ml model''' validate_model_config( @@ -55,7 +89,8 @@ def register_ml_model(onnx_model: ModelProto, model = MlModel(name=model_config.name, upload=str(upload_id), metadata=model_config.metadata, display_name=model_config.display_name, description=model_config.description) - ml_api.add_ml_model(model, _request_timeout=register_timeout) + res_name = ml_api.add_ml_model(model, _request_timeout=register_timeout) + return MlModelName.from_response(res_name) def validate_model_config(onnx_model: ModelProto, *, diff --git a/geoengine/permissions.py b/geoengine/permissions.py index 2e5697a..99f308a 100644 --- a/geoengine/permissions.py +++ b/geoengine/permissions.py @@ -6,19 +6,24 @@ from enum import Enum import ast -from typing import Dict, Literal, Any +from typing import Dict, List, Literal, Any, Union from uuid import UUID import geoengine_openapi_client +import geoengine_openapi_client.api +import geoengine_openapi_client.models +import geoengine_openapi_client.models.role from geoengine.auth import get_session from geoengine.datasets import DatasetName from geoengine.error import GeoEngineException from geoengine.layers import LayerCollectionId, LayerId +from geoengine.ml import MlModelName class RoleId: '''A wrapper for a role id''' + __role_id: UUID def __init__(self, role_id: UUID) -> None: self.__role_id = role_id @@ -48,6 +53,48 @@ def __repr__(self) -> str: return repr(self.__role_id) +class Role: + '''A wrapper for a role''' + name: str + id: RoleId + + def __init__(self, role_id: Union[UUID, RoleId, str], role_name: str): + ''' Create a role with name and id ''' + + if isinstance(role_id, UUID): + real_id = RoleId(role_id) + elif isinstance(role_id, str): + real_id = RoleId(UUID(role_id)) + else: + real_id = role_id + + self.id = real_id + self.name = role_name + + @classmethod + def from_response(cls, response: geoengine_openapi_client.models.role.Role) -> Role: + '''Parse a http response to an `RoleId`''' + + role_id = response.id + role_name = response.name + + return Role(role_id, role_name) + + def __eq__(self, other) -> bool: + '''Checks if two role ids are equal''' + if not isinstance(other, self.__class__): + return False + + return self.id == other.id and self.name == other.name + + def role_id(self) -> RoleId: + '''get the role id''' + return self.id + + def __repr__(self) -> str: + return 'id: ' + repr(self.id) + ', name: ' + repr(self.name) + + class UserId: '''A wrapper for a role id''' @@ -82,11 +129,14 @@ def __repr__(self) -> str: class Resource: '''A wrapper for a resource id''' - def __init__(self, resource_type: Literal['dataset', 'layer', 'layerCollection'], + id: str + type: Literal['dataset', 'layer', 'layerCollection', 'mlModel', 'project'] + + def __init__(self, resource_type: Literal['dataset', 'layer', 'layerCollection', 'mlModel', 'project'], resource_id: str) -> None: '''Create a resource id''' - self.__type = resource_type - self.__id = resource_id + self.type = resource_type + self.id = resource_id @classmethod def from_layer_id(cls, layer_id: LayerId) -> Resource: @@ -99,25 +149,100 @@ def from_layer_collection_id(cls, layer_collection_id: LayerCollectionId) -> Res return Resource('layerCollection', str(layer_collection_id)) @classmethod - def from_dataset_name(cls, dataset_name: DatasetName) -> Resource: - '''Create a resource id from a dataset id''' - return Resource('dataset', str(dataset_name)) + def from_dataset_name(cls, dataset_name: Union[DatasetName, str]) -> Resource: + '''Create a resource id from a dataset name''' + if isinstance(dataset_name, DatasetName): + dataset_name = str(dataset_name) + return Resource('dataset', dataset_name) + + @classmethod + def from_ml_model_name(cls, ml_model_name: Union[MlModelName, str]) -> Resource: + '''Create a resource from an ml model name''' + if isinstance(ml_model_name, MlModelName): + ml_model_name = str(ml_model_name) + return Resource('mlModel', ml_model_name) def to_api_dict(self) -> geoengine_openapi_client.Resource: '''Convert to a dict for the API''' inner: Any = None - if self.__type == "layer": - inner = geoengine_openapi_client.LayerResource(type="layer", id=self.__id) - elif self.__type == "layerCollection": - inner = geoengine_openapi_client.LayerCollectionResource(type="layerCollection", id=self.__id) - elif self.__type == "project": - inner = geoengine_openapi_client.ProjectResource(type="project", id=self.__id) - elif self.__type == "dataset": - inner = geoengine_openapi_client.DatasetResource(type="dataset", id=self.__id) + if self.type == "layer": + inner = geoengine_openapi_client.LayerResource(type="layer", id=self.id) + elif self.type == "layerCollection": + inner = geoengine_openapi_client.LayerCollectionResource(type="layerCollection", id=self.id) + elif self.type == "project": + inner = geoengine_openapi_client.ProjectResource(type="project", id=self.id) + elif self.type == "dataset": + inner = geoengine_openapi_client.DatasetResource(type="dataset", id=self.id) + elif self.type == "mlModel": + inner = geoengine_openapi_client.MlModelResource(type="mlModel", id=self.id) + else: + raise KeyError(f"Unknown resource type: {self.type}") return geoengine_openapi_client.Resource(inner) + @classmethod + def from_response(cls, response: geoengine_openapi_client.Resource) -> Resource: + '''Convert to a dict for the API''' + inner: Resource + if isinstance(response.actual_instance, geoengine_openapi_client.LayerResource): + inner = Resource('layer', response.actual_instance.id) + elif isinstance(response.actual_instance, geoengine_openapi_client.LayerCollectionResource): + inner = Resource('layerCollection', response.actual_instance.id) + elif isinstance(response.actual_instance, geoengine_openapi_client.ProjectResource): + inner = Resource('project', response.actual_instance.id) + elif isinstance(response.actual_instance, geoengine_openapi_client.DatasetResource): + inner = Resource('dataset', response.actual_instance.id) + elif isinstance(response.actual_instance, geoengine_openapi_client.MlModelResource): + inner = Resource('mlModel', response.actual_instance.id) + else: + raise KeyError(f"Unknown resource type from API: {response.actual_instance}") + return inner + + def __repr__(self): + return 'id: ' + repr(self.id) + ', type: ' + repr(self.type) + + def __eq__(self, value): + '''Checks if two listings are equal''' + if not isinstance(value, self.__class__): + return False + return self.id == value.id and self.type == value.type + + +class PermissionListing: + """ + PermissionListing + """ + permission: Permission + resource: Resource + role: Role + + def __init__(self, permission: Permission, resource: Resource, role: Role): + ''' Create a PermissionListing ''' + self.permission = permission + self.resource = resource + self.role = role + + @classmethod + def from_response(cls, response: geoengine_openapi_client.models.PermissionListing) -> PermissionListing: + ''' Transforms a response PermissionListing to a PermissionListing ''' + return PermissionListing( + permission=Permission.from_response(response.permission), + resource=Resource.from_response(response.resource), + role=Role.from_response(response.role) + ) + + def __eq__(self, other) -> bool: + '''Checks if two listings are equal''' + if not isinstance(other, self.__class__): + return False + return self.permission == other.permission and self.resource == other.resource and self.role == other.role + + def __repr__(self) -> str: + return 'Role: ' + repr(self.role) + ', ' \ + + 'Resource: ' + repr(self.resource) + ', ' \ + + 'Permission: ' + repr(self.permission) + class Permission(str, Enum): '''A permission''' @@ -128,6 +253,10 @@ def to_api_dict(self) -> geoengine_openapi_client.Permission: '''Convert to a dict for the API''' return geoengine_openapi_client.Permission(self.value) + @classmethod + def from_response(cls, response: geoengine_openapi_client.Permission) -> Permission: + return Permission(response) + ADMIN_ROLE_ID: RoleId = RoleId(UUID("d5328854-6190-4af9-ad69-4e74b0961ac9")) REGISTERED_USER_ROLE_ID: RoleId = RoleId(UUID("4e8081b6-8aa6-4275-af0c-2fa2da557d28")) @@ -164,6 +293,24 @@ def remove_permission(role: RoleId, resource: Resource, permission: Permission, )) +def list_permissions(resource: Resource, timeout: int = 60, offset=0, limit=20) -> List[PermissionListing]: + '''Lists the roles and permissions assigned to a ressource''' + + session = get_session() + + with geoengine_openapi_client.ApiClient(session.configuration) as api_client: + permission_api = geoengine_openapi_client.PermissionsApi(api_client) + res = permission_api.get_resource_permissions_handler( + resource_id=resource.id, + resource_type=resource.type, + offset=offset, + limit=limit, + _request_timeout=timeout + ) + + return [PermissionListing.from_response(r) for r in res] + + def add_role(name: str, timeout: int = 60) -> RoleId: """Add a new role. Requires admin role.""" diff --git a/setup.cfg b/setup.cfg index d6b0af4..327dded 100644 --- a/setup.cfg +++ b/setup.cfg @@ -18,7 +18,7 @@ package_dir = packages = find: python_requires = >=3.9 install_requires = - geoengine-openapi-client == 0.0.18 + geoengine-openapi-client == 0.0.19 geopandas >=0.9,<0.15 matplotlib >=3.5,<3.8 numpy >=1.21,<2.1 diff --git a/tests/test_ml.py b/tests/test_ml.py index 5fe0308..0d7e9d2 100644 --- a/tests/test_ml.py +++ b/tests/test_ml.py @@ -6,7 +6,7 @@ import numpy as np from geoengine_openapi_client.models import MlModelMetadata, RasterDataType import geoengine as ge -from . import UrllibMocker +from tests.ge_test import GeoEngineTestInstance class WorkflowStorageTests(unittest.TestCase): @@ -24,43 +24,19 @@ def test_uploading_onnx_model(self): onnx_clf = to_onnx(clf, training_x[:1], options={'zipmap': False}, target_opset=9) - with UrllibMocker() as m: - session_id = "c4983c3e-9b53-47ae-bda9-382223bd5081" - request_headers = {'Authorization': f'Bearer {session_id}'} - - m.post('http://mock-instance/anonymous', json={ - "id": session_id, - "project": None, - "view": None - }) - - upload_id = "c314ff6d-3e37-41b4-b9b2-3669f13f7369" - - m.post('http://mock-instance/upload', json={ - "id": upload_id - }, request_headers=request_headers) - - m.post('http://mock-instance/ml/models', - expected_request_body={ - "description": "A simple decision tree model", - "displayName": "Decision Tree", - "metadata": { - "fileName": "model.onnx", - "inputType": "F32", - "numInputBands": 2, - "outputType": "I64" - }, - "name": "foo", - "upload": upload_id - }, - request_headers=request_headers) - - ge.initialize("http://mock-instance") - - ge.register_ml_model( + # TODO: use `enterContext(cm)` instead of `with cm:` in Python 3.11 + with GeoEngineTestInstance() as ge_instance: + ge_instance.wait_for_ready() + + ge.initialize(ge_instance.address()) + + session = ge.get_session() + model_name = f"{session.user_id}:foo" + + res_name = ge.register_ml_model( onnx_model=onnx_clf, model_config=ge.ml.MlModelConfig( - name="foo", + name=model_name, metadata=MlModelMetadata( file_name="model.onnx", input_type=RasterDataType.F32, @@ -71,12 +47,33 @@ def test_uploading_onnx_model(self): description="A simple decision tree model", ) ) + self.assertEqual(str(res_name), model_name) + + # Now test permission setting and removal + ge.add_permission( + ge.REGISTERED_USER_ROLE_ID, ge.Resource.from_ml_model_name(res_name), ge.Permission.READ + ) + + expected = ge.permissions.PermissionListing( + permission=ge.Permission.READ, + resource=ge.Resource.from_ml_model_name(res_name), + role=ge.permissions.Role(ge.REGISTERED_USER_ROLE_ID, 'user') + ) + + self.assertIn(expected, ge.permissions.list_permissions(ge.Resource.from_ml_model_name(res_name))) + ge.remove_permission( + ge.REGISTERED_USER_ROLE_ID, ge.Resource.from_ml_model_name(res_name), ge.Permission.READ + ) + + self.assertNotIn(expected, ge.permissions.list_permissions(ge.Resource.from_ml_model_name(res_name))) + + # failing tests with self.assertRaises(ge.InputException) as exception: - ge.register_ml_model( + _res_name = ge.register_ml_model( onnx_model=onnx_clf, model_config=ge.ml.MlModelConfig( - name="foo", + name=model_name, metadata=MlModelMetadata( file_name="model.onnx", input_type=RasterDataType.F32, @@ -93,10 +90,10 @@ def test_uploading_onnx_model(self): ) with self.assertRaises(ge.InputException) as exception: - ge.register_ml_model( + _res_name = ge.register_ml_model( onnx_model=onnx_clf, model_config=ge.ml.MlModelConfig( - name="foo", + name=model_name, metadata=MlModelMetadata( file_name="model.onnx", input_type=RasterDataType.F64, diff --git a/tests/test_upload.py b/tests/test_upload.py index 8dd60af..ec75315 100644 --- a/tests/test_upload.py +++ b/tests/test_upload.py @@ -13,7 +13,7 @@ class UploadTests(unittest.TestCase): '''Test runner regarding upload functionality''' def setUp(self) -> None: - ge.reset(False) + ge.reset(logout=False) def test_upload(self): # TODO: use `enterContext(cm)` instead of `with cm:` in Python 3.11