Add a new treeDefault import feature

specifyweb/backend/trees/tree_mutations.py

@@ -0,0 +1,66 @@
+
+from specifyweb.backend.permissions.permissions import PermissionTarget, PermissionTargetAction
+
+
+class TaxonMutationPT(PermissionTarget):
+    resource = "/tree/edit/taxon"
+    merge = PermissionTargetAction()
+    move = PermissionTargetAction()
+    synonymize = PermissionTargetAction()
+    desynonymize = PermissionTargetAction()
+    repair = PermissionTargetAction()
+
+
+class GeographyMutationPT(PermissionTarget):
+    resource = "/tree/edit/geography"
+    merge = PermissionTargetAction()
+    move = PermissionTargetAction()
+    synonymize = PermissionTargetAction()
+    desynonymize = PermissionTargetAction()
+    repair = PermissionTargetAction()
+
+
+class StorageMutationPT(PermissionTarget):
+    resource = "/tree/edit/storage"
+    merge = PermissionTargetAction()
+    move = PermissionTargetAction()
+    bulk_move = PermissionTargetAction()
+    synonymize = PermissionTargetAction()
+    desynonymize = PermissionTargetAction()
+    repair = PermissionTargetAction()
+
+
+class GeologictimeperiodMutationPT(PermissionTarget):
+    resource = "/tree/edit/geologictimeperiod"
+    merge = PermissionTargetAction()
+    move = PermissionTargetAction()
+    synonymize = PermissionTargetAction()
+    desynonymize = PermissionTargetAction()
+    repair = PermissionTargetAction()
+
+
+class LithostratMutationPT(PermissionTarget):
+    resource = "/tree/edit/lithostrat"
+    merge = PermissionTargetAction()
+    move = PermissionTargetAction()
+    synonymize = PermissionTargetAction()
+    desynonymize = PermissionTargetAction()
+    repair = PermissionTargetAction()
+
+class TectonicunitMutationPT(PermissionTarget):
+    resource = "/tree/edit/tectonicunit"
+    merge = PermissionTargetAction()
+    move = PermissionTargetAction()
+    synonymize = PermissionTargetAction()
+    desynonymize = PermissionTargetAction()
+    repair = PermissionTargetAction()
+
+def perm_target(tree):
+    return {
+        'taxon': TaxonMutationPT,
+        'geography': GeographyMutationPT,
+        'storage': StorageMutationPT,
+        'geologictimeperiod': GeologictimeperiodMutationPT,
+        'lithostrat': LithostratMutationPT,
+        'tectonicunit':TectonicunitMutationPT
+    }[tree]

specifyweb/backend/trees/urls.py

@@ -20,4 +20,9 @@
         re_path(r'^(?P<treedef>\d+)/(?P<parentid>\w+)/(?P<sortfield>\w+)/$', views.tree_view),
         path('repair/', views.repair_tree),
     ])),
+
+    # Create new trees
+    path('create_default_tree/', views.create_default_tree_view),
+    re_path(r'^create_default_tree/status/(?P<task_id>[^/]+)/$', views.default_tree_upload_status),
+    re_path(r'^create_default_tree/abort/(?P<task_id>[^/]+)/$', views.abort_default_tree_creation),
 ]

specifyweb/backend/trees/utils.py

@@ -1,9 +1,21 @@
-from typing import Tuple, List
+from typing import Any, Callable, Dict, Iterator, Optional
+import json
+import requests
+import csv
+import time
+from requests.exceptions import ChunkedEncodingError, ConnectionError
+
+from django.db import transaction
 from django.db.models import Q, Count, Model
 
+from specifyweb.backend.notifications.models import Message
+from specifyweb.celery_tasks import LogErrorsTask, app
 import specifyweb.specify.models as spmodels
 from specifyweb.specify.datamodel import datamodel, Table
 
+import logging
+logger = logging.getLogger(__name__)
+
 lookup = lambda tree: (tree.lower() + 'treedef')
 
 SPECIFY_TREES = {"taxon", "storage", "geography", "geologictimeperiod", "lithostrat", 'tectonicunit'}
@@ -84,4 +96,260 @@ def get_models(name: str):
     tree_rank_model = get_treedefitem_model(name)
     tree_node_model = getattr(spmodels, name.lower().title())
 
-    return tree_def_model, tree_rank_model, tree_node_model
+    return tree_def_model, tree_rank_model, tree_node_model
+
+def initialize_default_tree(tree_type: str, discipline, tree_name: str, rank_names_lst: list):
+    """Creates an initial empty tree."""
+    with transaction.atomic():
+        tree_def_model, tree_rank_model, tree_node_model = get_models(tree_type)
+
+        # Uniquify name
+        tree_def = None
+        unique_tree_name = tree_name
+        if tree_def_model.objects.filter(name=tree_name).exists():
+            i = 1
+            while tree_def_model.objects.filter(name=f"{tree_name}_{i}").exists():
+                i += 1
+            unique_tree_name = f"{tree_name}_{i}"
+
+        # Create tree definition
+        tree_def, _ = tree_def_model.objects.get_or_create(
+            name=unique_tree_name,
+            discipline=discipline
+        )
+
+        # Create tree ranks
+        treedefitems_bulk = []
+        rank_id = 0
+        for i, rank_name in enumerate(rank_names_lst):
+            treedefitems_bulk.append(
+                tree_rank_model(
+                    treedef=tree_def,
+                    name=rank_name,
+                    title=rank_name.capitalize(), # TODO: allow rank name configuration
+                    rankid=int(rank_id),
+                    isenforced=(rank_id == 0),
+                    isinfullname=(i >= len(rank_names_lst)-3)
+                )
+            )
+            rank_id += 10
+        if treedefitems_bulk:
+            tree_rank_model.objects.bulk_create(treedefitems_bulk, ignore_conflicts=False)
+
+            # Create root node
+            # TODO: Avoid having duplicated code from add_root endpoint
+            root_rank = tree_rank_model.objects.get(treedef=tree_def, rankid=0)
+            tree_node, _ = tree_node_model.objects.get_or_create(
+                name="Root",
+                fullname="Root",
+                nodenumber=1,
+                definition=tree_def,
+                definitionitem=root_rank,
+                parent=None
+            )
+
+        tree_name = tree_def.name
+        return tree_name
+
+def add_default_tree_record(tree_type: str, discipline, row: dict, tree_name: str, tree_cfg: dict[str, Any]):
+    """
+    Given one CSV row, the discipline, and a rank configuration dictionary,
+    walk through the 'ranks' in order, creating or updating each tree record and linking
+    it to its parent.
+    """
+    tree_def_model, tree_rank_model, tree_node_model = get_models(tree_type)
+    tree_def = tree_def_model.objects.get(name=tree_name)
+    parent = tree_node_model.objects.get(name='Root', fullname='Root', definition=tree_def)
+    rank_id = 10
+
+    for rank_map in tree_cfg['ranks']:
+        rank = next(iter(rank_map))
+        fields_map = rank_map[rank]
+
+        value = row.get(rank)
+        if not value:
+            continue
+
+        defaults = {}
+        for csv_col, model_field in fields_map.items():
+            if csv_col == rank:
+                continue
+            v = row.get(csv_col)
+            if v:
+                defaults[model_field] = v
+
+        # Get the rank by the column name.
+        # It should already exist by this point, but worst case it will be generated here.
+        treedef_item, _ = tree_rank_model.objects.get_or_create(
+            name=rank,
+            treedef=tree_def,
+            defaults={
+                'title': rank.capitalize(),
+                'rankid': rank_id
+            }
+        )
+
+        obj = tree_node_model.objects.filter(
+            name=value,
+            fullname=value,
+            definition=tree_def,
+            definitionitem=treedef_item,
+            parent=parent,
+        ).first()
+        if obj is None:
+            data = {
+                'name': value,
+                'fullname': value,
+                'definition': tree_def,
+                'definitionitem': treedef_item,
+                'parent': parent,
+                'rankid': treedef_item.rankid,
+                **defaults
+            }
+            obj = tree_node_model(**data)
+            obj.save(skip_tree_extras=True)
+
+        # if not taxon_obj and defaults:
+        #     for f, v in defaults.items():
+        #         setattr(taxon_obj, f, v)
+        #     taxon_obj.save()
+
+        parent = obj
+        rank_id += 10
+
+@app.task(base=LogErrorsTask, bind=True)
+def create_default_tree_task(self, url: str, discipline_id: int, tree_discipline_name: str, specify_collection_id: int,
+                             specify_user_id: int, tree_cfg: dict, row_count: Optional[int], tree_name: str):
+    logger.info(f'starting task {str(self.request.id)}')
+
+    specify_user = spmodels.Specifyuser.objects.get(id=specify_user_id)
+    discipline = spmodels.Discipline.objects.get(id=discipline_id)
+
+    Message.objects.create(
+        user=specify_user,
+        content=json.dumps({
+            'type': 'create-default-tree-running',
+            'name': tree_name,
+            'taskid': str(self.request.id),
+            'collection_id': specify_collection_id,
+        })
+    )
+
+    current = 0
+    total = 1
+    def progress(cur: int, additional_total: int=0) -> None:
+        nonlocal current, total
+        current += cur
+        total += additional_total
+        if current > total:
+            current = total
+        self.update_state(state='RUNNING', meta={'current': current, 'total': total})
+
+    def set_tree(name: str) -> None:
+        # Final tree name after being made unique.
+        nonlocal tree_name
+        tree_name = name
+
+    try:
+        tree_type = 'taxon'
+        if tree_discipline_name in SPECIFY_TREES:
+            # non-taxon tree
+            tree_type = tree_discipline_name
+
+        rank_count = len(tree_cfg['ranks'])
+
+        total_rows = 0
+        if row_count:
+            total_rows = row_count-2
+        progress(0, total_rows)
+        with transaction.atomic():
+            for row in stream_csv_from_url(url, discipline, rank_count, tree_type, tree_name, set_tree):
+                add_default_tree_record(tree_type, discipline, row, tree_name, tree_cfg)
+                progress(1, 0)
+    except Exception as e:
+        Message.objects.create(
+            user=specify_user,
+            content=json.dumps({
+                'type': 'create-default-tree-failed',
+                'name': tree_name,
+                'taskid': str(self.request.id),
+                'collection_id': specify_collection_id,
+                # 'error': str(e)
+            })
+        )
+        raise
+
+    Message.objects.create(
+        user=specify_user,
+        content=json.dumps({
+            'type': 'create-default-tree-completed',
+            'name': tree_name,
+            'taskid': str(self.request.id),
+            'collection_id': specify_collection_id,
+        })
+    )
+
+def stream_csv_from_url(url: str, discipline, rank_count: int, tree_type: str, initial_tree_name: str, set_tree: Callable[[str], None]) -> Iterator[Dict[str, str]]:
+    """
+    Streams a taxon CSV from a URL. Yields each row.
+    """
+    chunk_size = 8192
+    max_retries = 5
+
+    def lines_iter() -> Iterator[str]:
+        # Streams data from the server in -chunks-, yields -lines-.
+        buffer = b""
+        bytes_downloaded = 0
+        retries = 0
+
+        headers = {}
+        while True:
+            # Request data starting from the last downloaded bytes
+            if bytes_downloaded > 0:
+                headers['Range'] = f'bytes={bytes_downloaded}-'
+
+            try:
+                with requests.get(url, stream=True, timeout=(5, 30), headers=headers) as resp:
+                    resp.raise_for_status()
+                    for chunk in resp.iter_content(chunk_size=chunk_size):
+                        chunk_length = len(chunk)
+                        if chunk_length == 0:
+                            continue
+                        buffer += chunk
+                        bytes_downloaded += chunk_length
+
+                        # Extract all lines from chunk
+                        while True:
+                            new_line_index = buffer.find(b'\n')
+                            if new_line_index == -1: break
+                            line = buffer[:new_line_index + 1] # extract line
+                            buffer = buffer[new_line_index + 1 :] # clear read buffer
+                            yield line.decode('utf-8-sig', errors='replace')
+
+                    if buffer:
+                        # yield last line
+                        yield buffer.decode('utf-8-sig', errors='replace')
+                    return
+            except (ChunkedEncodingError, ConnectionError) as e:
+                # Trigger retry
+                if retries < max_retries:
+                    retries += 1
+                    time.sleep(0.5*retries)
+                    continue
+                raise
+            except Exception:
+                raise
+
+    reader = csv.DictReader(lines_iter())
+
+    rank_names_lst = reader.fieldnames[:rank_count]
+    rank_names_lst.insert(0, "Root") # Add Root rank
+
+    logger.debug(f"Creating default tree with the following {rank_count} ranks:")
+    logger.debug(rank_names_lst)
+
+    tree_name = initialize_default_tree(tree_type, discipline, initial_tree_name, rank_names_lst)
+    set_tree(tree_name)
+
+    for row in reader:
+        yield row