9508: Fix place tree cursor

gramps/gen/db/generic.py

@@ -287,22 +287,6 @@ class Cursor:
     def close(self):
         pass
 
-class TreeCursor(Cursor):
-
-    def __init__(self, db, map):
-        """
-        """
-        self.db = db
-        Cursor.__init__(self, map)
-
-    def __iter__(self):
-        """
-        Iterator
-        """
-        handles = self.db.get_place_handles(sort_handles=True)
-        for handle in handles:
-            yield (handle, self.db.get_raw_place_data(handle))
-
 class DbGeneric(DbWriteBase, DbReadBase, UpdateCallback, Callback):
     """
     A Gramps Database Backend. This replicates the grampsdb functions.
@@ -1454,8 +1438,8 @@ class DbGeneric(DbWriteBase, DbReadBase, UpdateCallback, Callback):
     def get_place_cursor(self):
         return Cursor(self._iter_raw_place_data)
 
-    def get_place_tree_cursor(self, *args, **kwargs):
+    def get_place_tree_cursor(self):
-        return TreeCursor(self, self.place_map)
+        return Cursor(self._iter_raw_place_tree_data)
 
     def get_person_cursor(self):
         return Cursor(self._iter_raw_person_data)
@@ -2564,3 +2548,12 @@ class DbGeneric(DbWriteBase, DbReadBase, UpdateCallback, Callback):
                     if surname_obj:
                         surname = surname_obj.surname
         return (given_name, surname)
+
+    def _get_place_data(self, place):
+        """
+        Given a Place, return the first PlaceRef handle.
+        """
+        enclosed_by = ""
+        for placeref in place.get_placeref_list():
+            enclosed_by = placeref.ref
+        return enclosed_by

gramps/plugins/db/dbapi/dbapi.py

@@ -226,6 +226,7 @@ class DBAPI(DbGeneric):
         self.dbapi.execute('CREATE TABLE place '
                            '('
                            'handle VARCHAR(50) PRIMARY KEY NOT NULL, '
+                           'enclosed_by VARCHAR(50), '
                            'order_by TEXT, '
                            'gramps_id TEXT, '
                            'blob_data BLOB'
@@ -296,6 +297,8 @@ class DBAPI(DbGeneric):
                            'ON media(gramps_id)')
         self.dbapi.execute('CREATE INDEX place_order_by '
                            'ON place(order_by)')
+        self.dbapi.execute('CREATE INDEX place_enclosed_by '
+                           'ON place(enclosed_by)')
         self.dbapi.execute('CREATE INDEX place_gramps_id '
                            'ON place(gramps_id)')
         self.dbapi.execute('CREATE INDEX tag_order_by '
@@ -800,6 +803,20 @@ class DBAPI(DbGeneric):
                     yield (row[0].encode('utf8'), pickle.loads(row[1]))
                 rows = cursor.fetchmany()
 
+    def _iter_raw_place_tree_data(self):
+        """
+        Return an iterator over raw data in the place hierarchy.
+        """
+        to_do = ['']
+        sql = 'SELECT handle, blob_data FROM place WHERE enclosed_by = ?'
+        while to_do:
+            handle = to_do.pop()
+            self.dbapi.execute(sql, [handle])
+            rows = self.dbapi.fetchall()
+            for row in rows:
+                to_do.append(row[0])
+                yield (row[0].encode('utf8'), pickle.loads(row[1]))
+
     def reindex_reference_map(self, callback):
         """
         Reindex all primary records in the database.
@@ -1078,6 +1095,9 @@ class DBAPI(DbGeneric):
             sets.append("order_by = ?")
             values.append(self._order_by_person_key(obj))
         if table == 'Place':
+            handle = self._get_place_data(obj)
+            sets.append("enclosed_by = ?")
+            values.append(handle)
             sets.append("order_by = ?")
             values.append(self._order_by_place_key(obj))
         if table == 'Source':