Fix use after cursor closed error

Author: omad

cubedash/index/api.py

@@ -1,17 +1,18 @@
 from abc import ABC, abstractmethod
 from collections.abc import Generator, Iterable, Sequence
+from contextlib import contextmanager
 from datetime import date, datetime, timedelta
 from typing import Any
 from uuid import UUID
 
 from datacube.index import Index
 from datacube.model import Dataset, MetadataType, Product, Range
 from datacube.model.fields import Field
-from sqlalchemy import Result, Row, Select
+from sqlalchemy import DDL, Result, Row, Select
 from sqlalchemy.sql import ColumnElement
 from sqlalchemy.sql.elements import ClauseElement, Label
 
-from cubedash.summary._schema import PleaseRefresh
+from cubedash.summary._schema import CUBEDASH_SCHEMA, PleaseRefresh
 
 
 class EmptyDbError(Exception):
@@ -30,8 +31,28 @@ def __init__(self, name: str, index: Index) -> None:
 
     # need to add an odc_index accessor
     def execute_query(self, query):
+        """Execute query and return all rows"""
         with self.engine.begin() as conn:
-            return conn.execute(query)
+            return conn.execute(query).fetchall()
+
+    @contextmanager
+    def execute_query_lazy(self, query):
+        """
+        Execute potentially large query, keeping the connection alive
+
+        Example usage::
+
+            with explorer_index.execute_query_lazy(query) as result:
+                for row in result:
+        """
+        with self.engine.connect() as conn:
+            result = conn.execute(query)
+            yield result
+
+    def drop_all(self):
+        """Drop all cubedash-specific tables/schema."""
+        with self.engine.begin() as conn:
+            conn.execute(DDL(f"drop schema if exists {CUBEDASH_SCHEMA} cascade"))
 
     def make_dataset(self, row):
         # pylint: disable=protected-access

cubedash/summary/_stores.py

@@ -27,7 +27,7 @@
 from pygeofilter.parsers.cql2_text import parse as parse_cql2_text
 from shapely.geometry import MultiPolygon
 from shapely.geometry.base import BaseGeometry
-from sqlalchemy import DDL, Row, RowMapping, String, func, select
+from sqlalchemy import Row, RowMapping, String, func, select
 from sqlalchemy.dialects import postgresql as postgres
 from sqlalchemy.dialects.postgresql import TSTZRANGE
 from sqlalchemy.sql import Select
@@ -44,7 +44,7 @@
 from cubedash.index import EmptyDbError, ExplorerIndex
 from cubedash.index.postgis import ExplorerPgisIndex
 from cubedash.index.postgres import ExplorerPgIndex
-from cubedash.summary import RegionInfo, TimePeriodOverview, _extents, _schema
+from cubedash.summary import RegionInfo, TimePeriodOverview, _extents
 from cubedash.summary._extents import ProductArrival, RegionSummary
 from cubedash.summary._schema import PleaseRefresh
 from cubedash.summary._summarise import DEFAULT_TIMEZONE, Summariser
@@ -609,9 +609,7 @@ def drop_all(self) -> None:
         """
         Drop all cubedash-specific tables/schema.
         """
-        self.e_index.execute_query(
-            DDL(f"drop schema if exists {_schema.CUBEDASH_SCHEMA} cascade")
-        )
+        self.e_index.drop_all()
 
     def get(
         self,
@@ -1043,7 +1041,7 @@ def get_count(
             query = self._add_filter_to_query(
                 query, field_exprs, filter_lang, filter_cql
             )
-        result = self.e_index.execute_query(query).fetchall()
+        result = self.e_index.execute_query(query)
 
         if len(result) != 0:
             return result[0][0]
@@ -1123,21 +1121,26 @@ def search_items(
             offset
         )
 
-        for r in self.e_index.execute_query(query):
-            yield DatasetItem(
-                dataset_id=r.id,
-                bbox=_box2d_to_bbox(r.bbox) if r.bbox else None,
-                product_name=r.product_name,
-                geometry=(
-                    _get_shape(r.geometry, self.e_index.get_srid_name(r.geometry.srid))
-                    if r.geometry is not None
-                    else None
-                ),
-                region_code=r.region_code,
-                creation_time=r.creation_time,
-                center_time=r.center_time,
-                odc_dataset=(self.index.datasets.get(r.id) if full_dataset else None),
-            )
+        with self.e_index.execute_query_lazy(query) as result:
+            for r in result:
+                yield DatasetItem(
+                    dataset_id=r.id,
+                    bbox=_box2d_to_bbox(r.bbox) if r.bbox else None,
+                    product_name=r.product_name,
+                    geometry=(
+                        _get_shape(
+                            r.geometry, self.e_index.get_srid_name(r.geometry.srid)
+                        )
+                        if r.geometry is not None
+                        else None
+                    ),
+                    region_code=r.region_code,
+                    creation_time=r.creation_time,
+                    center_time=r.center_time,
+                    odc_dataset=(
+                        self.index.datasets.get(r.id) if full_dataset else None
+                    ),
+                )
 
     def search_collections(
         self,