managers.py

import logging

from django.contrib.gis.db import models
from django.contrib.gis.geos import GEOSGeometry
from django.db.models import Q
from django.utils.translation import gettext_lazy as _

from rest_framework import serializers

from stac_api.utils import fromisoformat
from stac_api.utils import geometry_from_bbox
from stac_api.validators import validate_geometry

logger = logging.getLogger(__name__)


class ItemQuerySet(models.QuerySet):

    def filter_by_bbox(self, bbox):
        '''Filter a querystring with a given bbox

        This function adds a bbox filter to the queryset.

        Args:
            queryset:
                A django queryset (https://docs.djangoproject.com/en/3.0/ref/models/querysets/)
            bbox:
                A string defining a spatial bbox (f.ex. 5.96, 45.82, 10.49, 47.81)

        Returns:
            The queryset with the added spatial filter

        Raises:
            ValidationError: When the bbox does not contain 4 values. Or when the polygon build
            from the bbox string is invalid.
        '''
        try:
            logger.debug('Query parameter bbox = %s', bbox)
            bbox_geometry = geometry_from_bbox(bbox)
            validate_geometry(bbox_geometry)
        except (ValueError, serializers.ValidationError, IndexError) as error:
            logger.error(
                'Invalid bbox query parameter: '
                'Could not transform bbox "%s" to a polygon; %s'
                'f.ex. bbox=5.96, 45.82, 10.49, 47.81',
                bbox,
                error
            )
            raise serializers.ValidationError(
                _('Invalid bbox query parameter, '
                  ' has to contain 4 values. f.ex. bbox=5.96,45.82,10.49,47.81'),
            ) from None

        return self.filter(geometry__intersects=bbox_geometry)

    def filter_by_datetime(self, date_time):
        '''Filter a queryset by datetime

        Args:
            queryset:
                 A django queryset (https://docs.djangoproject.com/en/3.0/ref/models/querysets/)
            date_time:
                A string

        Returns:
            The queryset filtered by date_time
        '''
        start, end = self._parse_datetime_query(date_time)
        if end is not None:
            return self._filter_by_datetime_range(start, end)
        return self.filter(properties_datetime=start)

    def _filter_by_datetime_range(self, start_datetime, end_datetime):
        '''Filter a queryset by datetime range

        Helper function of filter_by_datetime

        Args:
            queryset:
                A django queryset (https://docs.djangoproject.com/en/3.0/ref/models/querysets/)
            start_datetime:
                A string with the start datetime
            end_datetime:
                A string with the end datetime
        Returns:
            The queryset filtered by datetime range
        '''
        if start_datetime == '..':
            # open start range
            return self.filter(
                Q(properties_datetime__lte=end_datetime) |
                Q(properties_end_datetime__lte=end_datetime)
            )
        if end_datetime == '..':
            # open end range
            return self.filter(
                Q(properties_datetime__gte=start_datetime) |
                Q(properties_start_datetime__gte=start_datetime)
            )
            # else fixed range
        return self.filter(
            Q(properties_datetime__range=(start_datetime, end_datetime)) | (
                Q(properties_start_datetime__gte=start_datetime) &
                Q(properties_end_datetime__lte=end_datetime)
            )
        )

    def filter_by_forecast_reference_datetime(self, date_time):
        '''Filter a queryset by forecast reference datetime

        Args:
            date_time:
                A string containing datetime like "2020-10-28T13:05:10Z"

        Returns:
            The queryset filtered by date_time
        '''
        start, end = self._parse_datetime_query(date_time)
        if end is not None:
            return self._filter_by_forecast_reference_datetime_range(start, end)
        return self.filter(forecast_reference_datetime=start)

    def _filter_by_forecast_reference_datetime_range(self, start_datetime, end_datetime):
        '''Filter a queryset by forecast reference datetime range

        Helper function of filter_by_forecast_reference_datetime

        Args:
            start_datetime:
                A string with the start datetime or ".." to denote open start
            end_datetime:
                A string with the end datetime or ".." to denote open end
        Returns:
            The queryset filtered by datetime range
        '''
        if start_datetime == '..':
            # open start range
            return self.filter(forecast_reference_datetime__lte=end_datetime)
        if end_datetime == '..':
            # open end range
            return self.filter(forecast_reference_datetime__gte=start_datetime)
            # else fixed range
        return self.filter(forecast_reference_datetime__range=(start_datetime, end_datetime))

    def _parse_datetime_query(self, date_time):
        '''Parse the datetime query as specified in the api-spec.md.

        A helper function of filter_by_datetime

        Args:
            date_time: string
                Datetime as string (should be in isoformat)

        Returns:
            queryset filtered by datetime

        Raises:
            ValidationError: When the date_time string is not a valid isoformat
        '''
        start, sep, end = date_time.partition('/')
        if start == '':
            start = '..'
        try:
            if start != '..':
                start = fromisoformat(start)
            if end and end != '..':
                end = fromisoformat(end)
        except ValueError as error:
            logger.error(
                'Invalid datetime query parameter "%s", must be isoformat; %s', date_time, error
            )
            raise serializers.ValidationError(
                _('Invalid datetime query parameter, must be isoformat')
            ) from None

        if end == '':
            end = None

        if start == '..' and (end is None or end == '..'):
            logger.error(
                'Invalid datetime query parameter "%s"; '
                'cannot start with open range when no end range is defined',
                date_time
            )
            raise serializers.ValidationError(
                _('Invalid datetime query parameter, '
                  'cannot start with open range when no end range is defined')
            )
        return start, end

    def filter_by_item_name(self, item_names_array):
        '''Filter by item names parameter

        Args:
            item_names_array: list[string]
                An array of ids (string)

        Returns:
            queryset filtered by a list of ids
        '''
        return self.filter(name__in=item_names_array)

    def filter_by_collections(self, collections_array):
        '''Filter by collections parameter

        Args:
            collections_array: list[string]
                An array of collections (string)

        Returns:
            queryset filtered by a list of collections
        '''
        return self.filter(collection__name__in=collections_array)

    def filter_by_intersects(self, intersects):
        '''Filter by intersects parameter

        Args:
            intersects: string
                Is a geojson formatted string

        Returns:
            queryset filtered by intersects

        Raises:
            ValueError or GDALException: When the Geojson is not a valid geometry
        '''
        the_geom = GEOSGeometry(intersects)
        return self.filter(geometry__intersects=the_geom)

    def filter_by_forecast_horizon(self, duration):
        '''Filter by forecast horizon

        Args:
            duration: string
                ISO 8601 duration string

        Returns:
            queryset filtered by forecast horizon
        '''
        return self.filter(forecast_horizon=duration)

    def filter_by_forecast_duration(self, duration):
        '''Filter by forecast duration

        Args:
            duration: string
                ISO 8601 duration string

        Returns:
            queryset filtered by forecast duration
        '''
        return self.filter(forecast_duration=duration)

    def filter_by_forecast_variable(self, val):
        '''Filter by forecast variable

        Args:
            val: string
                foracast variable value

        Returns:
            queryset filtered by forecast variable
        '''
        return self.filter(forecast_variable=val)

    def filter_by_forecast_perturbed(self, pert):
        '''Filter by forecast perturbed

        Args:
            pert: boolean
                foracast perturbed

        Returns:
            queryset filtered by forecast perturbed
        '''
        return self.filter(forecast_perturbed=pert)

    def filter_by_query(self, query):
        '''Filter by the query parameter

        Args:
            query: dict
                {"attribute": {"operator": "value"}}
        Returns:
            queryset filtered by query
        '''
        for attribute in query:
            for operator in query[attribute]:
                value = query[attribute][operator]  # get the values given by the operator

                if attribute in ["updated", "created", "cf:standard_name", "unit"]:
                    prefix = ""
                else:
                    prefix = "properties_"

                model_field = attribute.replace(":", "_")

                # __eq does not exist, but = does it as well
                if operator == 'eq':
                    query_filter = f"{prefix}{model_field}"
                else:
                    query_filter = f"{prefix}{model_field}__{operator.lower()}"
                # PB-2354: This is a bug: As it is now, we cannot query by multiple fields
                # because we return already for the first operator.
                #
                # See this test to demonstrate the bug:
                #
                #    tests.tests_10.test_search_endpoint.SearchEndpointTestCF.test_multiple_cf
                #
                # We keep this is as it is as noone seemed to use that feature so far.
                # Instead, we intend to fix it properly by replacing the Query Extension by the
                # Filter Extension, which is the recommended solution anyway.
                return self.filter(**{query_filter: value})


class ItemManager(models.Manager):

    def get_queryset(self):
        # pylint: disable=not-callable
        return ItemQuerySet(self.model, using=self._db).select_related('collection')

    def filter_by_bbox(self, bbox):
        return self.get_queryset().filter_by_bbox(bbox)

    def filter_by_datetime(self, date_time):
        return self.get_queryset().filter_by_datetime(date_time)

    def filter_by_collections(self, collections_array):
        return self.get_queryset().filter_by_collections(collections_array)

    def filter_by_item_name(self, item_name_array):
        return self.get_queryset().filter_by_item_name(item_name_array)

    def filter_by_query(self, query):
        return self.get_queryset().filter_by_query(query)

    def filter_by_forecast_reference_datetime(self, date_time):
        return self.get_queryset().filter_by_forecast_reference_datetime(date_time)

    def filter_by_forecast_horizon(self, duration):
        return self.get_queryset().filter_by_forecast_horizon(duration)

    def filter_by_forecast_duration(self, duration):
        return self.get_queryset().filter_by_forecast_duration(duration)

    def filter_by_forecast_variable(self, val):
        return self.get_queryset().filter_by_forecast_variable(val)

    def filter_by_forecast_perturbed(self, pert):
        return self.get_queryset().filter_by_forecast_perturbed(pert)


class AssetUploadQuerySet(models.QuerySet):

    def filter_by_status(self, status):
        '''Filter a query with a given status

        Args:
            status:
                A string defining the status

        Returns:
            The queryset with the added status filter
        '''

        return self.filter(status=status)


class AssetUploadManager(models.Manager):

    def get_queryset(self):
        # pylint: disable=not-callable
        return AssetUploadQuerySet(self.model, using=self._db).select_related('asset')

    def filter_by_status(self, status):
        return self.get_queryset().filter_by_status(status)