Merge pull request #209 from geo-engine/get_or_create_unique_collection

Get or create unique collection

Author: ChristianBeilschmidt

geoengine/__init__.py

@@ -3,40 +3,40 @@
 from pkg_resources import get_distribution
 from requests import utils
 from pydantic import ValidationError
-
 from geoengine_openapi_client.exceptions import BadRequestException, OpenApiException, ApiTypeError, ApiValueError, \
     ApiKeyError, ApiAttributeError, ApiException, NotFoundException
 from geoengine_openapi_client import UsageSummaryGranularity
-from .auth import Session, get_session, initialize, reset
-from .colorizer import Colorizer, ColorBreakpoint, LinearGradientColorizer, PaletteColorizer, \
-    LogarithmicGradientColorizer
-from .datasets import upload_dataframe, StoredDataset, add_dataset, volumes, AddDatasetProperties, \
-    delete_dataset, list_datasets, DatasetListOrder, OgrSourceDatasetTimeType, OgrOnError
-from .error import GeoEngineException, InputException, UninitializedException, TypeException, \
-    MethodNotCalledOnPlotException, MethodNotCalledOnRasterException, MethodNotCalledOnVectorException, \
-    SpatialReferenceMismatchException, check_response_for_error, ModificationNotOnLayerDbException, \
-    InvalidUrlException, MissingFieldInResponseException, OGCXMLError
-from .layers import Layer, LayerCollection, LayerListing, LayerCollectionListing, \
-    LayerId, LayerCollectionId, LayerProviderId, \
-    layer_collection, layer
-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
+
+from . import workflow_builder
+from .raster_workflow_rio_writer import RasterWorkflowRioWriter
+from .raster import RasterTile2D
+from .workflow import WorkflowId, Workflow, workflow_by_id, register_workflow, get_quota, update_quota, data_usage, \
+    data_usage_summary
+from .util import clamp_datetime_ms_ns
+from .resource_identifier import LAYER_DB_PROVIDER_ID, LAYER_DB_ROOT_COLLECTION_ID, DatasetName, UploadId, \
+    LayerId, LayerCollectionId, LayerProviderId, Resource, MlModelName
 from .types import QueryRectangle, GeoTransform, \
     RasterResultDescriptor, Provenance, UnitlessMeasurement, ContinuousMeasurement, \
     ClassificationMeasurement, BoundingBox2D, TimeInterval, SpatialResolution, SpatialPartition2D, \
     RasterSymbology, VectorSymbology, VectorDataType, VectorResultDescriptor, VectorColumnInfo, \
     FeatureDataType, RasterBandDescriptor, DEFAULT_ISO_TIME_FORMAT, RasterColorizer, SingleBandRasterColorizer, \
     MultiBandRasterColorizer
-
-from .util import clamp_datetime_ms_ns
-from .workflow import WorkflowId, Workflow, workflow_by_id, register_workflow, get_quota, update_quota, data_usage, \
-    data_usage_summary
-from .raster import RasterTile2D
-from .raster_workflow_rio_writer import RasterWorkflowRioWriter
-
-from . import workflow_builder
+from .tasks import Task, TaskId
+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, UserId, RoleId
+from .ml import register_ml_model, MlModelConfig
+from .layers import Layer, LayerCollection, LayerListing, LayerCollectionListing, \
+    layer_collection, layer
+from .error import GeoEngineException, InputException, UninitializedException, TypeException, \
+    MethodNotCalledOnPlotException, MethodNotCalledOnRasterException, MethodNotCalledOnVectorException, \
+    SpatialReferenceMismatchException, check_response_for_error, ModificationNotOnLayerDbException, \
+    InvalidUrlException, MissingFieldInResponseException, OGCXMLError
+from .auth import Session, get_session, initialize, reset
+from .colorizer import Colorizer, ColorBreakpoint, LinearGradientColorizer, PaletteColorizer, \
+    LogarithmicGradientColorizer
+from .datasets import upload_dataframe, StoredDataset, add_dataset, volumes, AddDatasetProperties, \
+    delete_dataset, list_datasets, DatasetListOrder, OgrSourceDatasetTimeType, OgrOnError, \
+    add_or_replace_dataset_with_permissions, dataset_info_by_name
 
 
 DEFAULT_USER_AGENT = f'geoengine-python/{get_distribution("geoengine").version}'

geoengine/datasets.py

@@ -5,20 +5,24 @@
 from __future__ import annotations
 from abc import abstractmethod
 from pathlib import Path
-from typing import List, NamedTuple, Optional, Union, Literal
+from typing import List, NamedTuple, Optional, Union, Literal, Tuple
 from enum import Enum
 from uuid import UUID
 import tempfile
 from attr import dataclass
+import geoengine_openapi_client
+import geoengine_openapi_client.exceptions
+import geoengine_openapi_client.models
 import numpy as np
 import geopandas as gpd
-import geoengine_openapi_client
 from geoengine import api
 from geoengine.error import InputException, MissingFieldInResponseException
 from geoengine.auth import get_session
 from geoengine.types import Provenance, RasterSymbology, TimeStep, \
     TimeStepGranularity, VectorDataType, VectorResultDescriptor, VectorColumnInfo, \
     UnitlessMeasurement, FeatureDataType
+from geoengine.resource_identifier import Resource, UploadId, DatasetName
+from geoengine.permissions import RoleId, Permission, add_permission
 
 
 class UnixTimeStampType(Enum):
@@ -256,71 +260,6 @@ def to_api_enum(self) -> geoengine_openapi_client.OgrSourceErrorSpec:
         return geoengine_openapi_client.OgrSourceErrorSpec(self.value)
 
 
-class DatasetName:
-    '''A wrapper for a dataset name'''
-
-    __dataset_name: str
-
-    def __init__(self, dataset_name: str) -> None:
-        self.__dataset_name = dataset_name
-
-    @classmethod
-    def from_response(cls, response: geoengine_openapi_client.CreateDatasetHandler200Response) -> DatasetName:
-        '''Parse a http response to an `DatasetName`'''
-        return DatasetName(response.dataset_name)
-
-    def __str__(self) -> str:
-        return self.__dataset_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.__dataset_name == other.__dataset_name  # pylint: disable=protected-access
-
-    def to_api_dict(self) -> geoengine_openapi_client.CreateDatasetHandler200Response:
-        return geoengine_openapi_client.CreateDatasetHandler200Response(
-            dataset_name=str(self.__dataset_name)
-        )
-
-
-class UploadId:
-    '''A wrapper for an upload id'''
-
-    __upload_id: UUID
-
-    def __init__(self, upload_id: UUID) -> None:
-        self.__upload_id = upload_id
-
-    @classmethod
-    def from_response(cls, response: geoengine_openapi_client.AddCollection200Response) -> UploadId:
-        '''Parse a http response to an `UploadId`'''
-        return UploadId(UUID(response.id))
-
-    def __str__(self) -> str:
-        return str(self.__upload_id)
-
-    def __repr__(self) -> str:
-        return str(self)
-
-    def __eq__(self, other) -> bool:
-        '''Checks if two upload ids are equal'''
-        if not isinstance(other, self.__class__):
-            return False
-
-        return self.__upload_id == other.__upload_id  # pylint: disable=protected-access
-
-    def to_api_dict(self) -> geoengine_openapi_client.AddCollection200Response:
-        '''Converts the upload id to a dict for the api'''
-        return geoengine_openapi_client.AddCollection200Response(
-            id=str(self.__upload_id)
-        )
-
-
 class AddDatasetProperties():
     '''The properties for adding a dataset'''
     name: Optional[str]
@@ -601,6 +540,42 @@ def add_dataset(data_store: Union[Volume, UploadId],
     return DatasetName.from_response(response)
 
 
+def add_or_replace_dataset_with_permissions(data_store: Union[Volume, UploadId],
+                                            properties: AddDatasetProperties,
+                                            meta_data: geoengine_openapi_client.MetaDataDefinition,
+                                            permission_tuples: Optional[List[Tuple[RoleId, Permission]]] = None,
+                                            replace_existing=False,
+                                            timeout: int = 60) -> DatasetName:
+    '''
+    Add a dataset to the Geo Engine and set permissions.
+    Replaces existing datasets if forced!
+    '''
+    # pylint: disable=too-many-arguments,too-many-positional-arguments
+
+    def add_dataset_and_permissions() -> DatasetName:
+        dataset_name = add_dataset(data_store=data_store, properties=properties, meta_data=meta_data, timeout=timeout)
+        if permission_tuples is not None:
+            dataset_res = Resource.from_dataset_name(dataset_name)
+            for (role, perm) in permission_tuples:
+                add_permission(role, dataset_res, perm, timeout=timeout)
+        return dataset_name
+
+    if properties.name is None:
+        dataset_name = add_dataset_and_permissions()
+
+    else:
+        dataset_name = DatasetName(properties.name)
+        dataset_info = dataset_info_by_name(dataset_name)
+        if dataset_info is None:  # dataset is not existing
+            dataset_name = add_dataset_and_permissions()
+        else:
+            if replace_existing:  # dataset exists and we overwrite it
+                delete_dataset(dataset_name)
+                dataset_name = add_dataset_and_permissions()
+
+    return dataset_name
+
+
 def delete_dataset(dataset_name: DatasetName, timeout: int = 60) -> None:
     '''Delete a dataset. The dataset must be owned by the caller.'''
 
@@ -636,3 +611,25 @@ def list_datasets(offset: int = 0,
         )
 
     return response
+
+
+def dataset_info_by_name(
+        dataset_name: Union[DatasetName, str], timeout: int = 60
+) -> geoengine_openapi_client.models.Dataset | None:
+    '''Get dataset information.'''
+
+    if not isinstance(dataset_name, DatasetName):
+        dataset_name = DatasetName(dataset_name)
+
+    session = get_session()
+
+    with geoengine_openapi_client.ApiClient(session.configuration) as api_client:
+        datasets_api = geoengine_openapi_client.DatasetsApi(api_client)
+        res = None
+        try:
+            res = datasets_api.get_dataset_handler(str(dataset_name), _request_timeout=timeout)
+        except geoengine_openapi_client.exceptions.BadRequestException as e:
+            e_body = e.body
+            if isinstance(e_body, str) and 'CannotLoadDataset' not in e_body:
+                raise e
+        return res

geoengine/layers.py

@@ -7,7 +7,7 @@
 from enum import auto
 from io import StringIO
 import os
-from typing import Any, Dict, Generic, List, Literal, NewType, Optional, TypeVar, Union, cast
+from typing import Any, Dict, Generic, List, Literal, Optional, TypeVar, Union, cast, Tuple
 from uuid import UUID
 import json
 from strenum import LowercaseStrEnum
@@ -16,15 +16,11 @@
 from geoengine.error import ModificationNotOnLayerDbException, InputException
 from geoengine.tasks import Task, TaskId
 from geoengine.types import Symbology
+from geoengine.permissions import RoleId, Permission, add_permission
 from geoengine.workflow import Workflow, WorkflowId
 from geoengine.workflow_builder.operators import Operator as WorkflowBuilderOperator
-
-LayerId = NewType('LayerId', str)
-LayerCollectionId = NewType('LayerCollectionId', str)
-LayerProviderId = NewType('LayerProviderId', UUID)
-
-LAYER_DB_PROVIDER_ID = LayerProviderId(UUID('ce5e84db-cbf9-48a2-9a32-d4b7cc56ea74'))
-LAYER_DB_ROOT_COLLECTION_ID = LayerCollectionId('05102bb3-a855-4a37-8a8a-30026a91fef1')
+from geoengine.resource_identifier import LayerCollectionId, LayerId, LayerProviderId, \
+    LAYER_DB_PROVIDER_ID, Resource
 
 
 class LayerCollectionListingType(LowercaseStrEnum):
@@ -319,6 +315,26 @@ def add_layer(self,
 
         return layer_id
 
+    def add_layer_with_permissions(self,
+                                   name: str,
+                                   description: str,
+                                   workflow: Union[Dict[str, Any], WorkflowBuilderOperator],  # TODO: improve type
+                                   symbology: Optional[Symbology],
+                                   permission_tuples: Optional[List[Tuple[RoleId, Permission]]] = None,
+                                   timeout: int = 60) -> LayerId:
+        '''
+        Add a layer to this collection and set permissions.
+        '''
+
+        layer_id = self.add_layer(name, description, workflow, symbology, timeout)
+
+        if permission_tuples is not None:
+            res = Resource.from_layer_id(layer_id)
+            for (role, perm) in permission_tuples:
+                add_permission(role, res, perm)
+
+        return layer_id
+
     def add_existing_layer(self,
                            existing_layer: Union[LayerListing, Layer, LayerId],
                            timeout: int = 60):
@@ -456,6 +472,63 @@ def search(self, search_string: str, *,
 
         return listings
 
+    def get_or_create_unique_collection(
+            self,
+            collection_name: str,
+            create_collection_description: Optional[str] = None,
+            delete_existing_with_same_name: bool = False,
+            create_permissions_tuples: Optional[List[Tuple[RoleId, Permission]]] = None
+    ) -> LayerCollection:
+        '''
+        Get a unique child by name OR if it does not exist create it.
+        Removes existing collections with same name if forced!
+        Sets permissions if the collection is created from a list of tuples
+        '''
+        parent_collection = self.reload()  # reload just to be safe since self's state change on the server
+        existing_collections = parent_collection.get_items_by_name(collection_name)
+
+        if delete_existing_with_same_name and len(existing_collections) > 0:
+            for c in existing_collections:
+                actual = c.load()
+                if isinstance(actual, LayerCollection):
+                    actual.remove()
+            parent_collection = parent_collection.reload()
+            existing_collections = parent_collection.get_items_by_name(collection_name)
+
+        if len(existing_collections) == 0:
+            new_desc = create_collection_description if create_collection_description is not None else collection_name
+            new_collection = parent_collection.add_collection(collection_name, new_desc)
+            new_ressource = Resource.from_layer_collection_id(new_collection)
+
+            if create_permissions_tuples is not None:
+                for (role, perm) in create_permissions_tuples:
+                    add_permission(role, new_ressource, perm)
+            parent_collection = parent_collection.reload()
+            existing_collections = parent_collection.get_items_by_name(collection_name)
+
+        if len(existing_collections) == 0:
+            raise KeyError(
+                f"No collection with name {collection_name} exists in {parent_collection.name} and none was created!"
+            )
+
+        if len(existing_collections) > 1:
+            raise KeyError(f"Multiple collections with name {collection_name} exist in {parent_collection.name}")
+
+        res = existing_collections[0].load()
+        if isinstance(res, Layer):
+            raise TypeError(f"Found a Layer not a Layer collection for {collection_name}")
+
+        return cast(LayerCollection, existing_collections[0].load())  # we know that it is a collection since check that
+
+    def __eq__(self, other):
+        ''' Tests if two layer listings are identical '''
+        if not isinstance(other, self.__class__):
+            return False
+
+        return self.name == other.name and self.description == other.description \
+            and self.provider_id == other.provider_id \
+            and self.collection_id == other.collection_id and self.items == other.items
+
 
 @dataclass(repr=False)
 class Layer:
@@ -663,7 +736,7 @@ def layer(layer_id: LayerId,
 
     with geoengine_openapi_client.ApiClient(session.configuration) as api_client:
         layers_api = geoengine_openapi_client.LayersApi(api_client)
-        response = layers_api.layer_handler(str(layer_provider_id), layer_id, _request_timeout=timeout)
+        response = layers_api.layer_handler(str(layer_provider_id), str(layer_id), _request_timeout=timeout)
 
     return Layer.from_response(response)