django-gisserver 1.5.0__py3-none-any.whl → 2.0__py3-none-any.whl

gisserver/output/iters.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable
+from itertools import islice
+from typing import TypeVar
+
+from django.db import connections, models
+from lru import LRU
+
+M = TypeVar("M", bound=models.Model)
+
+DEFAULT_SQL_CHUNK_SIZE = 2000  # allow unit tests to alter this.
+
+logger = logging.getLogger(__name__)
+
+
+class CountingIterator(Iterable[M]):
+    """A simple iterator that counts how many results are given."""
+
+    def __init__(self, iterator: Iterable[M], max_results=0):
+        self._iterator = iterator
+        self._number_returned = 0
+        self._in_iterator = False
+        self._max_results = max_results
+        self._has_more = None
+
+    def __iter__(self):
+        # Count the number of returned items while reading them.
+        # Tried using map(itemgetter(0), zip(model_iter, count_iter)) but that's not faster.
+        self._in_iterator = True
+        try:
+            self._number_returned = 0
+            for instance in self._iterator:
+                if self._max_results and self._number_returned == self._max_results:
+                    self._has_more = True
+                    break
+                self._number_returned += 1
+                yield instance
+        finally:
+            if self._max_results and self._has_more is None:
+                self._has_more = False  # ignored the sentinel item
+            self._in_iterator = False
+
+    @property
+    def number_returned(self) -> int:
+        """Tell how many objects the iterator processed"""
+        if self._in_iterator:
+            raise RuntimeError("Can't read number of returned results during iteration")
+        return self._number_returned
+
+    @property
+    def has_more(self) -> bool | None:
+        return self._has_more
+
+
+class lru_dict(dict):
+    """A 'defaultdict' with LRU items for each value."""
+
+    def __init__(self, max_size):
+        super().__init__()
+        self.max_size = max_size
+
+    def __missing__(self, key):
+        logger.debug("Creating cache for prefetches of '%s'", key)
+        value = LRU(self.max_size)
+        self[key] = value
+        return value
+
+
+class ChunkedQuerySetIterator(Iterable[M]):
+    """An optimal strategy to perform ``prefetch_related()`` on large datasets.
+
+    It fetches data from the queryset in chunks,
+    and performs ``prefetch_related()`` behavior on each chunk.
+
+    Django's ``QuerySet.prefetch_related()`` works by loading the whole queryset into memory,
+    and performing an analysis of the related objects to fetch. When working on large datasets,
+    this is very inefficient as more memory is consumed. Instead, ``QuerySet.iterator()``
+    is preferred here as it returns instances while reading them. Nothing is stored in memory.
+    Hence, both approaches are fundamentally incompatible. This class performs a
+    mixed strategy: load a chunk, and perform prefetches for that particular batch.
+
+    As extra performance benefit, a local cache avoids prefetching the same records
+    again when the next chunk is analysed. It has a "least recently used" cache to avoid
+    flooding the caches when foreign keys constantly point to different unique objects.
+    """
+
+    def __init__(self, queryset: models.QuerySet, chunk_size=None, sql_chunk_size=None):
+        """
+        :param queryset: The queryset to iterate over, that has ``prefetch_related()`` data.
+        :param chunk_size: The size of each segment to analyse in-memory for related objects.
+        :param sql_chunk_size: The size of each segment to fetch from the database,
+            used when server-side cursors are not available. The default follows Django behavior.
+        """
+        self.queryset = queryset
+        self.sql_chunk_size = sql_chunk_size or DEFAULT_SQL_CHUNK_SIZE
+        self.chunk_size = chunk_size or self.sql_chunk_size
+        self._fk_caches = lru_dict(self.chunk_size // 2)
+        self._number_returned = 0
+        self._in_iterator = False
+
+    def __iter__(self):
+        # Using iter() ensures the ModelIterable is resumed with the next chunk.
+        self._number_returned = 0
+        self._in_iterator = True
+        try:
+            qs_iter = iter(self._get_queryset_iterator())
+
+            # Keep fetching chunks
+            while True:
+                instances = list(islice(qs_iter, self.chunk_size))
+                if not instances:
+                    break
+
+                # Perform prefetches on this chunk:
+                if self.queryset._prefetch_related_lookups:
+                    self._add_prefetches(instances)
+
+                # And return to parent loop
+                yield from instances
+                self._number_returned += len(instances)
+        finally:
+            self._in_iterator = False
+
+    def _get_queryset_iterator(self) -> Iterable:
+        """The body of queryset.iterator(), while circumventing prefetching."""
+        # The old code did return `self.queryset.iterator(chunk_size=self.sql_chunk_size)`
+        # However, Django 4 supports using prefetch_related() with iterator() in that scenario.
+        #
+        # This code is the core of Django's QuerySet.iterator() that only produces the
+        # old-style iteration, without any prefetches. Those are added by this class instead.
+        use_chunked_fetch = not connections[self.queryset.db].settings_dict.get(
+            "DISABLE_SERVER_SIDE_CURSORS"
+        )
+        iterable = self.queryset._iterable_class(
+            self.queryset, chunked_fetch=use_chunked_fetch, chunk_size=self.sql_chunk_size
+        )
+
+        yield from iterable
+
+    @property
+    def number_returned(self) -> int:
+        """Tell how many objects the iterator processed"""
+        if self._in_iterator:
+            raise RuntimeError("Can't read number of returned results during iteration")
+        return self._number_returned
+
+    def _add_prefetches(self, instances: list[M]):
+        """Merge the prefetched objects for this batch with the model instances."""
+        if self._fk_caches:
+            # Make sure prefetch_related_objects() doesn't have
+            # to fetch items again that infrequently changes.
+            all_restored = self._restore_caches(instances)
+            if all_restored:
+                logger.debug("Restored all prefetches from cache")
+                return
+
+        # Reuse the Django machinery for retrieving missing sub objects.
+        # and analyse the ForeignKey caches to allow faster prefetches next time
+        logger.debug("Perform additional prefetches for %d objects", len(instances))
+        models.prefetch_related_objects(instances, *self.queryset._prefetch_related_lookups)
+        self._persist_prefetch_cache(instances)
+
+    def _persist_prefetch_cache(self, instances):
+        """Store the prefetched data so it can be applied to the next batch"""
+        for instance in instances:
+            for lookup, obj in instance._state.fields_cache.items():
+                if obj is not None:
+                    cache = self._fk_caches[lookup]
+                    cache[obj.pk] = obj
+
+    def _restore_caches(self, instances) -> bool:
+        """Restore prefetched data to the new set of instances.
+        This avoids unneeded prefetching of the same ForeignKey relation.
+        """
+        if not instances:
+            return True
+        if not self._fk_caches:
+            return False
+
+        all_restored = True
+
+        for lookup, cache in self._fk_caches.items():
+            field = instances[0]._meta.get_field(lookup)
+            if not hasattr(field, "attname"):
+                logger.debug(
+                    "Unable to restore prefetches for '%s' (%s)", lookup, field.__class__.__name__
+                )
+                # Retrieving prefetches from ForeignObjectRel wouldn't work here.
+                # Let standard prefetch_related() take over.
+                all_restored = False
+                continue
+
+            logger.debug("Restoring prefetches for '%s'", lookup)
+            for instance in instances:
+                id_value = getattr(instance, field.attname)
+                if id_value is None:
+                    continue
+
+                obj = cache.get(id_value, None)
+                if obj is not None:
+                    instance._state.fields_cache[lookup] = obj
+                else:
+                    all_restored = False
+
+        return all_restored

gisserver/output/results.py

@@ -20,10 +20,13 @@ from gisserver import conf
 from gisserver.features import FeatureType
 from gisserver.geometries import BoundingBox
 
-from .utils import ChunkedQuerySetIterator, CountingIterator
+from .iters import ChunkedQuerySetIterator, CountingIterator
 
 if typing.TYPE_CHECKING:
-    from gisserver.queries import FeatureProjection, QueryExpression
+    from gisserver.projection import FeatureProjection, QueryExpression
+
+
+CALCULATE = -9999999
 
 
 class SimpleFeatureCollection:
@@ -36,21 +39,29 @@ class SimpleFeatureCollection:
     def __init__(
         self,
         source_query: QueryExpression,
-        feature_type: FeatureType,
+        feature_types: list[FeatureType],
         queryset: models.QuerySet,
         start: int,
         stop: int,
+        number_matched: int | None = CALCULATE,
     ):
         self.source_query = source_query
-        self.feature_type = feature_type
+        self.feature_types = feature_types
         self.queryset = queryset
         self.start = start
         self.stop = stop
+        self._number_matched = number_matched
 
         self._result_cache = None
         self._result_iterator = None
         self._has_more = None
 
+        # Tell that is a resultType=hits request.
+        # Typically, start and stop are 0. However, for resultType=hits with count,
+        # that does not apply. Instead, it detects whether the known amount is already provided.
+        # Detecting that queryset.none() is provided won't work, as that can be used by IdOperator too.
+        self._is_hits_request = number_matched is not None and number_matched != CALCULATE
+
     def __iter__(self) -> Iterable[models.Model]:
         """Iterate through all results.
 
@@ -58,7 +69,7 @@ class SimpleFeatureCollection:
         the results can either be cached first, or be streamed without caching.
         This picks the best-performance scenario in most cases.
         """
-        if self.start == self.stop == 0:
+        if self._is_hits_request:
             self._result_cache = []
 
         if self._result_cache is not None:
@@ -83,7 +94,7 @@ class SimpleFeatureCollection:
         if self._result_cache is not None:
             # In case the results were read already, reuse that.
             return iter(self._result_cache)
-        elif self.start == self.stop == 0:
+        elif self._is_hits_request:
             # resulttype=hits
             return iter([])
         else:
@@ -129,8 +140,7 @@ class SimpleFeatureCollection:
         if self._result_iterator is not None:
             raise RuntimeError("Results for feature collection are read twice.")
 
-        if self.start == self.stop == 0:
-            # resulttype=hits
+        if self._is_hits_request:
             self._result_cache = []
         else:
             # This still allows prefetch_related() to work,
@@ -171,8 +181,8 @@ class SimpleFeatureCollection:
     @cached_property
     def number_returned(self) -> int:
         """Return the number of results for this page."""
-        if self.start == self.stop == 0:
-            return 0  # resulttype=hits
+        if self._is_hits_request:
+            return 0
         elif self._result_iterator is not None:
             # When requesting the data after the fact, results are counted.
             return self._result_iterator.number_returned
@@ -183,9 +193,19 @@ class SimpleFeatureCollection:
             self.fetch_results()
             return len(self._result_cache)
 
-    @cached_property
+    @property
     def number_matched(self) -> int:
         """Return the total number of matches across all pages."""
+        if self._is_hits_request:
+            if self.stop:
+                # resulttype=hits&COUNT=n should minimize how many are "matched".
+                return min(self._number_matched, self.stop - self.start)
+            else:
+                return self._number_matched
+        elif self._number_matched != CALCULATE:
+            # Return previously cached result
+            return self._number_matched
+
         if self._is_surely_last_page:
             # For resulttype=results, an expensive COUNT query can be avoided
             # when this is the first and only page or the last page.
@@ -197,17 +217,24 @@ class SimpleFeatureCollection:
             # Otherwise, it becomes SELECT COUNT(*) FROM (SELECT AsGML(..), ...)
             key: value
             for key, value in qs.query.annotations.items()
-            if not key.startswith("_as_")
+            if not key.startswith("_as_") and not key.startswith("_As")  # AsGML / AsEWKT
         }
         if clean_annotations != qs.query.annotations:
             qs = self.queryset.all()  # make a clone to allow editing
             qs.query.annotations = clean_annotations
 
-        return qs.count()
+        # Calculate, cache and return
+        self._number_matched = qs.count()
+        return self._number_matched
 
     @property
     def _is_surely_last_page(self):
         """Return true when it's totally clear this is the last page."""
+        if self.start == self.stop == 0:
+            return True  # hits request without count
+        elif self._is_hits_request:
+            return False
+
         # Optimization to avoid making COUNT() queries when we can already know the answer.
         if self.stop == math.inf:
             return True  # Infinite page requested
@@ -226,27 +253,30 @@ class SimpleFeatureCollection:
         # For GeoJSON output, the iterator was read first, and `number_returned` is already filled in.
         # For GML output, the pagination details are requested first, and will fetch all data.
         # Hence, reading `number_returned` here can be quite an intensive operation.
-        page_size = self.stop - self.start  # is 0 for resulttype=hits
+        page_size = self.stop - self.start
         return page_size and (self.number_returned < page_size or self._has_more is False)
 
     @property
     def has_next(self):
-        if self.stop == math.inf:
+        if self.stop == math.inf or (self.start == self.stop == 0):
             return False
         elif self._has_more is not None:
             return self._has_more  # did page+1 record check, answer is known.
         elif self._is_surely_last_page:
             return False  # Fewer results than expected, answer is known.
 
-        # This will perform an slow COUNT() query...
-        return self.stop < self.number_matched
+        if self._is_hits_request:
+            return self.stop <= self._number_matched
+        else:
+            # This will perform an slow COUNT() query...
+            return self.stop < self.number_matched
 
     @cached_property
     def projection(self) -> FeatureProjection:
         """Provide the projection to render these results with."""
         # Note this attribute would technically be part of the 'query' object,
-        # but since the projection needs to be mapped per feature type, it's stored here for convenience.
-        return self.source_query.get_projection(self.feature_type)
+        # but since the projection needs to be calculated once, it's stored here for convenience.
+        return self.source_query.get_projection()
 
     def get_bounding_box(self) -> BoundingBox:
         """Determine bounding box of all items."""
@@ -255,11 +285,10 @@ class SimpleFeatureCollection:
         # Start with an obviously invalid bbox,
         # which corrects at the first extend_to_geometry call.
         bbox = BoundingBox(math.inf, math.inf, -math.inf, -math.inf)
-        geometry_field = self.feature_type.resolve_element(
-            self.feature_type.geometry_field.name
-        ).child
+
+        # Allow the geometry to exist in a dotted relationship.
         for instance in self:
-            geometry_value = geometry_field.get_value(instance)
+            geometry_value = self.projection.get_main_geometry_value(instance)
             if geometry_value is None:
                 continue
 
@@ -268,9 +297,6 @@ class SimpleFeatureCollection:
         return bbox
 
 
-CALCULATE = -9999999
-
-
 class FeatureCollection:
     """WFS object that holds the result type for GetFeature.
     This object type is defined in the WFS spec.

gisserver/output/stored.py

@@ -0,0 +1,143 @@
+"""Outputting XML for the stored query logic."""
+
+from __future__ import annotations
+
+from io import StringIO
+from xml.etree.ElementTree import Element, tostring
+
+from gisserver.extensions.queries import QueryExpressionText, StoredQueryDescription
+from gisserver.output.utils import attr_escape, tag_escape
+from gisserver.parsers.xml import xmlns
+
+from .base import XmlOutputRenderer
+
+
+class ListStoredQueriesRenderer(XmlOutputRenderer):
+    """Rendering for the ``<wfs:ListStoredQueriesResponse>``."""
+
+    # XML Namespaces to include by default
+    xml_namespaces = {
+        xmlns.wfs20: "",
+        xmlns.xs: "xs",
+        xmlns.xsi: "xsi",
+    }
+
+    def __init__(self, operation, query_descriptions: list[StoredQueryDescription]):
+        """Take the list of stored queries to render."""
+        super().__init__(operation)
+        self.all_feature_types = operation.view.get_bound_feature_types()
+        self.query_descriptions = query_descriptions
+
+    def render_stream(self):
+        self.output = StringIO()
+        self.output.write(
+            '<?xml version="1.0" encoding="UTF-8"?>\n'
+            f"<ListStoredQueriesResponse"
+            f" {self.render_xmlns_attributes()}"
+            f' xsi:schemaLocation="http://www.opengis.net/wfs/2.0 http://schemas.opengis.net/wfs/2.0/wfs.xsd">\n'
+        )
+        for query_description in self.query_descriptions:
+            self.write_query(query_description)
+
+        self.output.write("</ListStoredQueriesResponse>\n")
+        return self.output.getvalue()
+
+    def write_query(self, query_description: StoredQueryDescription):
+        self.output.write(
+            f'  <StoredQuery id="{query_description.id}">\n'
+            f"    <Title>{tag_escape(query_description.title)}</Title>\n"
+        )
+
+        for expression in query_description.expressions:
+            return_types = expression.return_feature_types or self.all_feature_types
+            for return_type in return_types:
+                feature_qname = self.feature_to_qname(return_type)
+                self.output.write(
+                    f"    <ReturnFeatureType>{tag_escape(feature_qname)}</ReturnFeatureType>\n"
+                )
+
+        self.output.write("  </StoredQuery>\n")
+
+
+class DescribeStoredQueriesRenderer(XmlOutputRenderer):
+    """Rendering for the ``<wfs:DescribeStoredQueriesResponse>``."""
+
+    # XML Namespaces to include by default
+    xml_namespaces = {
+        xmlns.wfs20: "",
+        xmlns.xs: "xs",
+        xmlns.xsi: "xsi",
+    }
+
+    def __init__(self, operation, query_descriptions: list[StoredQueryDescription]):
+        """Take the list of stored queries to render."""
+        super().__init__(operation)
+        self.all_feature_types = operation.view.get_bound_feature_types()
+        self.query_descriptions = query_descriptions
+
+    def render_stream(self):
+        self.output = StringIO()
+        self.output.write(
+            '<?xml version="1.0" encoding="UTF-8"?>\n'
+            f"<DescribeStoredQueriesResponse"
+            f" {self.render_xmlns_attributes()}"
+            f' xsi:schemaLocation="http://www.opengis.net/wfs/2.0 http://schemas.opengis.net/wfs/2.0/wfs.xsd">\n'
+        )
+
+        for query_description in self.query_descriptions:
+            self.write_description(query_description)
+
+        self.output.write("</DescribeStoredQueriesResponse>\n")
+        return self.output.getvalue()
+
+    def write_description(self, query_description: StoredQueryDescription):
+        """Write the stored query description."""
+        self.output.write(
+            f'<StoredQueryDescription id="{attr_escape(query_description.id)}">\n'
+            f"  <Title>{tag_escape(query_description.title)}</Title>\n"
+            f"  <Abstract>{tag_escape(query_description.abstract)}</Abstract>\n"
+        )
+
+        # Declare parameters
+        for name, xsd_type in query_description.parameters.items():
+            type_qname = self.to_qname(xsd_type)
+            self.output.write(f'  <Parameter name="{attr_escape(name)}" type="{type_qname}"/>\n')
+
+        # The QueryExpressionText nodes allow code per return type.
+        for expression in query_description.expressions:
+            self.render_expression(expression)
+
+        self.output.write("</StoredQueryDescription>\n")
+
+    def render_expression(self, expression: QueryExpressionText):
+        """Render the 'QueryExpressionText' node (no body content for now)."""
+        is_private = "true" if expression.is_private else "false"
+        if expression.return_feature_types is None:
+            # for GetFeatureById
+            types = " ".join(self.feature_to_qname(ft) for ft in self.all_feature_types)
+        else:
+            types = " ".join(
+                self.feature_to_qname(return_type)
+                for return_type in expression.return_feature_types
+            )
+
+        if expression.is_private or not expression.implementation_text:
+            implementation_text = ""
+        elif isinstance(expression.implementation_text, Element):
+            # XML serialization (will recreate namespaces)
+            default_namespace = next(
+                (ns for ns, prefix in self.app_namespaces.items() if prefix == ""), None
+            )
+            implementation_text = tostring(
+                expression.implementation_text,
+                xml_declaration=False,
+                default_namespace=default_namespace,
+            )
+        else:
+            # Some raw content (e.g. language="python")
+            implementation_text = tag_escape(expression.implementation_text)
+
+        self.output.write(
+            f'  <QueryExpressionText isPrivate="{is_private}" language="{expression.language}"'
+            f' returnFeatureTypes="{types}">{implementation_text}</QueryExpressionText>\n'
+        )