raster.py

import pyproj
from importer.publisher import DataPublisher
import json
import logging
from pathlib import Path
from subprocess import PIPE, Popen
from typing import List

from django.conf import settings
from django.db.models import Q
from geonode.base.models import ResourceBase
from geonode.layers.models import Dataset
from geonode.resource.enumerator import ExecutionRequestAction as exa
from geonode.resource.manager import resource_manager
from geonode.resource.models import ExecutionRequest
from importer.api.exception import ImportException
from importer.celery_tasks import ErrorBaseTaskClass, import_orchestrator
from importer.handlers.base import BaseHandler
from importer.handlers.geotiff.exceptions import InvalidGeoTiffException
from importer.handlers.utils import create_alternate, should_be_imported
from importer.models import ResourceHandlerInfo
from importer.orchestrator import orchestrator
from osgeo import gdal
from importer.celery_app import importer_app
from geonode.storage.manager import storage_manager
from geonode.assets.handlers import asset_handler_registry
from geonode.assets.models import Asset
from geonode.assets.utils import create_link

logger = logging.getLogger(__name__)

gdal.UseExceptions()


class BaseRasterFileHandler(BaseHandler):
    """
    Handler to import Raster files into GeoNode data db
    It must provide the task_lists required to comple the upload
    """

    @property
    def default_geometry_column_name(self):
        return "geometry"

    @property
    def supported_file_extension_config(self):
        return NotImplementedError

    @staticmethod
    def get_geoserver_store_name(default=None):
        """
        Method that return the base store name where to save the data in geoserver
        and a boolean to know if the store should be created.
        For raster, the store is created during the geoserver publishing
        so we dont want to created it upfront
        """
        return default, False

    @staticmethod
    def is_valid(files, user):
        """
        Define basic validation steps
        """
        result = Popen("gdal_translate --version", stdout=PIPE, stderr=PIPE, shell=True)
        _, stderr = result.communicate()
        if stderr:
            raise ImportException(stderr)
        return True

    @staticmethod
    def can_handle(_data) -> bool:
        """
        This endpoint will return True or False if with the info provided
        the handler is able to handle the file or not
        """
        return False

    @staticmethod
    def has_serializer(_data) -> bool:
        """
        This endpoint will return True or False if with the info provided
        the handler is able to handle the file or not
        """
        return False

    @staticmethod
    def can_do(action) -> bool:
        """
        This endpoint will return True or False if with the info provided
        the handler is able to handle the file or not
        """
        return action in BaseHandler.ACTIONS

    @staticmethod
    def create_error_log(exc, task_name, *args):
        """
        This function will handle the creation of the log error for each message.
        This is helpful and needed, so each handler can specify the log as needed
        """
        return f"Task: {task_name} raised an error during actions for layer: {args[-1]}: {exc}"

    @staticmethod
    def extract_params_from_data(_data, action=None):
        """
        Remove from the _data the params that needs to save into the executionRequest object
        all the other are returned
        """
        if action == exa.COPY.value:
            title = json.loads(_data.get("defaults"))
            return {"title": title.pop("title"), "store_spatial_file": True}, _data

        return {
            "skip_existing_layers": _data.pop("skip_existing_layers", "False"),
            "overwrite_existing_layer": _data.pop("overwrite_existing_layer", "False"),
            "store_spatial_file": _data.pop("store_spatial_files", "True"),
            "source": _data.pop("source", "upload"),
        }, _data

    @staticmethod
    def publish_resources(resources: List[str], catalog, store, workspace):
        """
        Given a list of strings (which rappresent the table on geoserver)
        Will publish the resorces on geoserver
        """
        for _resource in resources:
            try:
                catalog.create_coveragestore(
                    _resource.get("name"),
                    path=_resource.get("raster_path"),
                    layer_name=_resource.get("name"),
                    workspace=workspace,
                    overwrite=True,
                    upload_data=False,
                )
            except Exception as e:
                if (
                    f"Resource named {_resource.get('name')} already exists in store:"
                    in str(e)
                ):
                    continue
                raise e
        return True

    def overwrite_geoserver_resource(
        self, resource: List[str], catalog, store, workspace
    ):
        # we need to delete the resource before recreating it
        self._delete_resource(resource, catalog, workspace)
        self._delete_store(resource, catalog, workspace)
        return self.publish_resources([resource], catalog, store, workspace)

    def _delete_store(self, resource, catalog, workspace):
        store = None
        possible_layer_name = [
            resource.get("name"),
            resource.get("name").split(":")[-1],
            f"{workspace.name}:{resource.get('name')}",
        ]
        for el in possible_layer_name:
            store = catalog.get_store(el, workspace=workspace)
            if store:
                break
        if store:
            catalog.delete(store, purge="all", recurse=True)
        return store

    def _delete_resource(self, resource, catalog, workspace):
        res = None
        possible_layer_name = [
            resource.get("name"),
            resource.get("name").split(":")[-1],
            f"{workspace.name}:{resource.get('name')}",
        ]
        for el in possible_layer_name:
            res = catalog.get_resource(el, workspace=workspace)
            if res:
                break
        if res:
            catalog.delete(res, purge="all", recurse=True)

    @staticmethod
    def delete_resource(instance):
        # it should delete the image from the geoserver data dir
        # for now we can rely on the geonode delete behaviour
        # since the file is stored on local
        pass

    @staticmethod
    def perform_last_step(execution_id):
        BaseHandler.perform_last_step(execution_id=execution_id)

    def extract_resource_to_publish(
        self, files, action, layer_name, alternate, **kwargs
    ):
        if action == exa.COPY.value:
            nl = kwargs.get("new_file_location") or kwargs.get("kwargs", {}).get("new_file_location", {})
            raster_path = None
            if nl:
                files_list = nl.get("files") or []
                raster_path = files_list[0] if files_list else None
            return [
                {
                    "name": alternate,
                    "crs": ResourceBase.objects.filter(
                        Q(alternate__icontains=layer_name)
                        | Q(title__icontains=layer_name)
                    )
                    .first()
                    .srid,
                    "raster_path": raster_path
                }
            ]

        layers = gdal.Open(files.get("base_file"))
        if not layers:
            return []
        return [
            {
                "name": alternate or layer_name,
                "crs": (
                    self.identify_authority(layers) if layers.GetSpatialRef() else None
                ),
                "raster_path": files.get("base_file"),
            }
        ]

    def identify_authority(self, layer):
        try:
            layer_wkt = layer.GetSpatialRef().ExportToWkt()
            _name = "EPSG"
            _code = pyproj.CRS(layer_wkt).to_epsg(min_confidence=20)
            if _code is None:
                layer_proj4 = layer.GetSpatialRef().ExportToProj4()
                _code = pyproj.CRS(layer_proj4).to_epsg(min_confidence=20)
                if _code is None:
                    raise Exception(
                        "CRS authority code not found, fallback to default behaviour"
                    )
        except Exception:
            spatial_ref = layer.GetSpatialRef()
            spatial_ref.AutoIdentifyEPSG()
            _name = spatial_ref.GetAuthorityName(None) or spatial_ref.GetAttrValue(
                "AUTHORITY", 0
            )
            _code = (
                spatial_ref.GetAuthorityCode("PROJCS")
                or spatial_ref.GetAuthorityCode("GEOGCS")
                or spatial_ref.GetAttrValue("AUTHORITY", 1)
            )
        return f"{_name}:{_code}"

    def import_resource(self, files: dict, execution_id: str, **kwargs) -> str:
        """
        Main function to import the resource.
        Internally will call the steps required to import the
        data inside the geonode_data database
        """
        # for the moment we skip the dyanamic model creation
        logger.info("Total number of layers available: 1")
        _exec = self._get_execution_request_object(execution_id)
        _input = {**_exec.input_params, **{"total_layers": 1}}
        orchestrator.update_execution_request_status(
            execution_id=str(execution_id), input_params=_input
        )

        try:
            filename = Path(files.get("base_file")).stem
            # start looping on the layers available
            layer_name = self.fixup_name(filename)

            should_be_overwritten = _exec.input_params.get("overwrite_existing_layer")
            # should_be_imported check if the user+layername already exists or not
            if should_be_imported(
                layer_name,
                _exec.user,
                skip_existing_layer=_exec.input_params.get("skip_existing_layer"),
                overwrite_existing_layer=should_be_overwritten,
            ):
                workspace = DataPublisher(None).workspace
                user_datasets = Dataset.objects.filter(alternate=f"{workspace.name}:{layer_name}")

                dataset_exists = user_datasets.exists()

                if dataset_exists and should_be_overwritten:
                    layer_name, alternate = (
                        layer_name,
                        user_datasets.first().alternate.split(":")[-1],
                    )
                elif not dataset_exists:
                    alternate = layer_name
                else:
                    alternate = create_alternate(layer_name, execution_id)

                import_orchestrator.apply_async(
                    (
                        files,
                        execution_id,
                        str(self),
                        "importer.import_resource",
                        layer_name,
                        alternate,
                        exa.IMPORT.value,
                    )
                )
                return layer_name, alternate, execution_id

        except Exception as e:
            logger.error(e)
            raise e
        return

    def create_geonode_resource(
        self,
        layer_name: str,
        alternate: str,
        execution_id: str,
        resource_type: Dataset = Dataset,
        asset=None,
        custom={},
    ):
        """
        Base function to create the resource into geonode. Each handler can specify
        and handle the resource in a different way
        """
        saved_dataset = resource_type.objects.filter(alternate__icontains=alternate)

        _exec = self._get_execution_request_object(execution_id)

        workspace = getattr(
            settings,
            "DEFAULT_WORKSPACE",
            getattr(settings, "CASCADE_WORKSPACE", "geonode"),
        )

        _overwrite = _exec.input_params.get("overwrite_existing_layer", False)
        # if the layer exists, we just update the information of the dataset by
        # let it recreate the catalogue
        if not saved_dataset.exists() and _overwrite:
            logger.warning(
                f"The dataset required {alternate} does not exists, but an overwrite is required, the resource will be created"
            )

        saved_dataset = resource_manager.create(
            None,
            resource_type=resource_type,
            defaults=dict(
                name=alternate,
                workspace=workspace,
                subtype="raster",
                alternate=f"{workspace}:{alternate}",
                dirty_state=True,
                title=layer_name,
                owner=_exec.user,
                asset=asset,
            ),
            custom=custom,
        )

        saved_dataset.refresh_from_db()

        self.handle_xml_file(saved_dataset, _exec)
        self.handle_sld_file(saved_dataset, _exec)

        resource_manager.set_thumbnail(None, instance=saved_dataset)

        ResourceBase.objects.filter(alternate=alternate).update(dirty_state=False)

        saved_dataset.refresh_from_db()
        return saved_dataset

    def overwrite_geonode_resource(
        self,
        layer_name: str,
        alternate: str,
        execution_id: str,
        resource_type: Dataset = Dataset,
        asset=None,
        custom={},        
    ):
        
        _exec = self._get_execution_request_object(execution_id)

        dataset = resource_type.objects.filter(alternate__icontains=alternate, owner=_exec.user)

        _overwrite = _exec.input_params.get("overwrite_existing_layer", False)
        # if the layer exists, we just update the information of the dataset by
        # let it recreate the catalogue

        if dataset.exists() and _overwrite:
            dataset = dataset.first()

            dataset = resource_manager.update(dataset.uuid, instance=dataset)

            self.handle_xml_file(dataset, _exec)
            self.handle_sld_file(dataset, _exec)

            resource_manager.set_thumbnail(
                dataset.uuid, instance=dataset, overwrite=True
            )
            dataset.refresh_from_db()
            return dataset
        elif not dataset.exists() and _overwrite:
            logger.warning(
                f"The dataset required {alternate} does not exists, but an overwrite is required, the resource will be created"
            )
            return self.create_geonode_resource(
                layer_name, alternate, execution_id, resource_type, asset, custom
            )
        elif not dataset.exists() and not _overwrite:
            logger.warning(
                "The resource does not exists, please use 'create_geonode_resource' to create one"
            )
        return

    def handle_xml_file(self, saved_dataset: Dataset, _exec: ExecutionRequest):
        _path = _exec.input_params.get("files", {}).get("xml_file", "")
        resource_manager.update(
            None,
            instance=saved_dataset,
            xml_file=_path,
            metadata_uploaded=True if _path else False,
            vals={"dirty_state": True},
        )

    def handle_sld_file(self, saved_dataset: Dataset, _exec: ExecutionRequest):
        _path = _exec.input_params.get("files", {}).get("sld_file", "")
        resource_manager.exec(
            "set_style",
            None,
            instance=saved_dataset,
            sld_file=_exec.input_params.get("files", {}).get("sld_file", ""),
            sld_uploaded=True if _path else False,
            vals={"dirty_state": True},
        )

    def create_resourcehandlerinfo(
        self,
        handler_module_path: str,
        resource: Dataset,
        execution_id: ExecutionRequest,
        **kwargs,
    ):
        """
        Create relation between the GeonodeResource and the handler used
        to create/copy it
        """
        ResourceHandlerInfo.objects.create(
            handler_module_path=str(handler_module_path),
            resource=resource,
            execution_request=execution_id,
            kwargs=kwargs.get("kwargs", {}),
        )

    def overwrite_resourcehandlerinfo(
        self,
        handler_module_path: str,
        resource: Dataset,
        execution_id: ExecutionRequest,
        **kwargs,
    ):
        """
        Overwrite the ResourceHandlerInfo
        """
        if resource.resourcehandlerinfo_set.exists():
            resource.resourcehandlerinfo_set.update(
                handler_module_path=handler_module_path,
                resource=resource,
                execution_request=execution_id,
                kwargs=kwargs.get("kwargs", {}) or kwargs,
            )
            return
        return self.create_resourcehandlerinfo(
            handler_module_path, resource, execution_id, **kwargs
        )
    def _prepare_assets_for_copy(self, resource, kwargs):
        """
        Prepare assets for copying.
        It gets the cloned asset and identifies other assets to be linked.
        """
        _nl = kwargs.get("new_file_location") or kwargs.get("kwargs", {}).get("new_file_location", {})
        _asset = _nl.get("asset")
        _asset_id = _nl.get("asset_id")
        if not _asset and _asset_id:
            _asset = Asset.objects.filter(pk=_asset_id).first()

        assets_to_link = Asset.objects.filter(link__resource=resource).exclude(title="Original")
        return _asset, assets_to_link

    def copy_geonode_resource(
        self,
        alternate: str,
        resource: Dataset,
        _exec: ExecutionRequest,
        data_to_update: dict,
        new_alternate: str,
        **kwargs,
    ):
        cloned_asset, assets_to_link = self._prepare_assets_for_copy(resource, kwargs)

        new_resource = self.create_geonode_resource(
            layer_name=data_to_update.get("title"),
            alternate=new_alternate,
            execution_id=str(_exec.exec_id),
            asset=cloned_asset,
        )
        [create_link(new_resource, asset) for asset in assets_to_link]

        new_resource.refresh_from_db()
        return new_resource

    def _get_execution_request_object(self, execution_id: str):
        return ExecutionRequest.objects.filter(exec_id=execution_id).first()

    @staticmethod
    def copy_original_file(dataset):
        """
        Copy the original file into a new location
        """
        return storage_manager.copy(dataset)

    def _import_resource_rollback(self, exec_id, istance_name=None, *args, **kwargs):
        """
        In the raster, this step just generate the alternate, no real action
        are done on the database
        """
        pass

    def _publish_resource_rollback(self, exec_id, instance_name=None, *args, **kwargs):
        """
        We delete the resource from geoserver
        """
        logger.info(
            f"Rollback publishing step in progress for execid: {exec_id} resource published was: {instance_name}"
        )
        exec_object = orchestrator.get_execution_object(exec_id)
        handler_module_path = exec_object.input_params.get("handler_module_path")
        publisher = DataPublisher(handler_module_path=handler_module_path)
        publisher.delete_resource(instance_name)


@importer_app.task(
    base=ErrorBaseTaskClass,
    name="importer.copy_raster_file",
    queue="importer.copy_raster_file",
    max_retries=1,
    acks_late=False,
    ignore_result=False,
    task_track_started=True,
)
def copy_raster_file(
    exec_id, actual_step, layer_name, alternate, handler_module_path, action, **kwargs
):
    """
    Perform a copy of the original raster file"""

    original_dataset = ResourceBase.objects.filter(alternate=alternate)
    if not original_dataset.exists():
        raise InvalidGeoTiffException("Dataset required does not exists")

    original_dataset = original_dataset.first()

    # Ensure the dataset has at least one Asset associated
    filters = {"link__resource": original_dataset, "title": "Original"}
    if not Asset.objects.filter().exists():
        raise InvalidGeoTiffException(
            "The dataset does not have any original asset associated; cannot copy the dataset"
        )
    original_asset = Asset.objects.filter(**filters).last()

    # The original asset is cloned here, as it is required for the creation of the new cloned resource. Other associated assets will be linked to the new resource later in the process.
    cloned_asset = asset_handler_registry.get_handler(original_asset).clone(original_asset)
    new_file_location = {
        "files": cloned_asset.location if getattr(cloned_asset, "location", None) else [],
        "asset_id": cloned_asset.id,
    }
    if not new_file_location["files"]:
        raise InvalidGeoTiffException("Could not determine the location of the copied file")

    sanitized_title = BaseHandler().fixup_name(original_dataset.title)
    new_dataset_alternate = create_alternate(sanitized_title, exec_id)

    additional_kwargs = {
        "original_dataset_alternate": original_dataset.alternate,
        "new_dataset_alternate": new_dataset_alternate,
        "new_file_location": new_file_location,
    }

    task_params = (
        {},
        exec_id,
        handler_module_path,
        actual_step,
        layer_name,
        new_dataset_alternate,
        action,
    )

    import_orchestrator.apply_async(task_params, additional_kwargs)

    return "copy_raster", layer_name, alternate, exec_id