Treatments: efficient view to get upstream and dowstream - CSV export of tables (refs #3108)

1 files changed, 57 insertions, 49 deletions

diff --git a/archaeological_finds/models.py b/archaeological_finds/models.py
index d1f194727..ea68fa807 100644
--- a/archaeological_finds/models.py
+++ b/archaeological_finds/models.py
@@ -449,43 +449,28 @@ class Find(BaseHistorizedItem, ImageModel, OwnPerms, ShortMenuItem):
             bf.context_record.operation.get_reference(),
             self.index)
 
-    def upstream_treatments(self):
-        treatments = []
-        base_finds = [bf.pk for bf in self.base_finds.all()]
-        if self.upstream_treatment and \
-                self.upstream_treatment.pk not in treatments:
-            treatments.append(
-                (self.upstream_treatment.upstream.distinct(
-                ).order_by('label').all(), self.upstream_treatment))
-            for upstream in self.upstream_treatment.upstream.all():
-                if upstream.pk != self.pk and not [
-                        bf.pk for bf in upstream.base_finds.all()
-                        if bf.pk in base_finds]:
-                    continue
-                for items, treatment in upstream.upstream_treatments():
-                    if treatment.pk not in treatments:
-                        treatments.append((treatment.upstream.order_by(
-                            'label').all(), treatment))
+    def _get_treatments(self, model, rel='upstream'):
+        treatments, findtreats = [], []
+        for findtreat in model.objects.filter(find_id=self.pk
+                ).order_by('treatment_nb', 'treatment__start_date',
+                           'treatment__end_date'
+                ).distinct().all():
+            if findtreat.pk in findtreats:
+                continue
+            findtreats.append(findtreat.pk)
+            q = getattr(findtreat.treatment, rel).distinct().order_by(
+                'label')
+            treatments.append((q.all(), findtreat.treatment))
         return treatments
 
+    def upstream_treatments(self):
+        return self._get_treatments(FindUpstreamTreatments, 'upstream')
+
     def downstream_treatments(self):
-        treatments = []
-        base_finds = [bf.pk for bf in self.base_finds.all()]
-        if self.downstream_treatment and \
-                self.downstream_treatment.pk not in treatments:
-            treatments.append(
-                (self.downstream_treatment.downstream.distinct(
-                ).order_by('label').all(), self.downstream_treatment))
-            for downstream in self.downstream_treatment.downstream.all():
-                if downstream.pk != self.pk and not [
-                        bf.pk for bf in downstream.base_finds.all()
-                        if bf.pk in base_finds]:
-                    continue
-                for items, treatment in downstream.downstream_treatments():
-                    if treatment.pk not in treatments:
-                        treatments.append((treatment.downstream.order_by(
-                            'label').all(), treatment))
-        return treatments
+        return self._get_treatments(FindDownstreamTreatments, 'downstream')
+
+    def all_treatments(self):
+        return self.upstream_treatments() + self.downstream_treatments()
 
     def get_department(self):
         bf = self.get_first_base_find()
@@ -705,7 +690,7 @@ class Treatment(BaseHistorizedItem, OwnPerms):
     other_location = models.CharField(_(u"Other location"), max_length=200,
                                       blank=True, null=True)
     person = models.ForeignKey(
-        Person, verbose_name=_(u"Person"), blank=True, null=True,
+        Person, verbose_name=_(u"Doer"), blank=True, null=True,
         on_delete=models.SET_NULL, related_name='treatments')
     start_date = models.DateField(_(u"Start date"), blank=True, null=True)
     end_date = models.DateField(_(u"End date"), blank=True, null=True)
@@ -758,9 +743,9 @@ class Treatment(BaseHistorizedItem, OwnPerms):
                 basket.items.add(new)
 
 
-class AbsFindTreatments(object):
+class AbsFindTreatments(models.Model):
     find = models.ForeignKey(Find, verbose_name=_(u"Find"),
-                             related_name='treatments')
+                             related_name='%(class)s_related')
     treatment = models.ForeignKey(Treatment, verbose_name=_(u"Treatment"),
                                   primary_key=True)
     # primary_key is set to prevent django to ask for an id column
@@ -768,8 +753,18 @@ class AbsFindTreatments(object):
     treatment_nb = models.IntegerField(_(u"Order"))
     TABLE_COLS = ['treatment__treatment_type',
                   'treatment__start_date', 'treatment__end_date',
-                  'treatment__location', 'treatment__person',
-                  'treatment_nb']
+                  'treatment__location', 'treatment__container',
+                  'treatment__person', 'treatment_nb']
+    EXTRA_FULL_FIELDS_LABELS = {
+        'treatment__treatment_type': _(u"Treatment type"),
+        'treatment__start_date': _(u"Start date"),
+        'treatment__end_date': _(u"End date"),
+        'treatment__location': _(u"Location"),
+        'treatment__container': _(u"Container"),
+        'treatment__person': _(u"Doer"),
+        'treatment__upstream': _(u"Related finds"),
+        'treatment__downstream': _(u"Related finds"),
+    }
 
     class Meta:
         abstract = True
@@ -779,9 +774,9 @@ class AbsFindTreatments(object):
             self.find, self.treatment, self.treatment_nb)
 
 
-class FindDownstreamTreatments(AbsFindTreatments):
+class FindUpstreamTreatments(AbsFindTreatments):
     """
-    CREATE VIEW find_downtreatments_tree AS
+    CREATE VIEW find_uptreatments_tree AS
         WITH RECURSIVE rel_tree AS (
           SELECT id AS find_id, upstream_treatment_id, downstream_treatment_id,
               1 AS level,
@@ -799,24 +794,30 @@ class FindDownstreamTreatments(AbsFindTreatments):
         SELECT DISTINCT find_id, path_info, level
           FROM rel_tree ORDER BY find_id;
 
-    CREATE VIEW find_downtreatments AS
+    CREATE VIEW find_uptreatments AS
         SELECT DISTINCT find_id,
             path_info[nb] AS treatment_id, level - nb + 1 AS treatment_nb
           FROM (SELECT *, generate_subscripts(path_info, 1) AS nb
-                FROM find_downtreatments_tree) y
+                FROM find_uptreatments_tree) y
          WHERE path_info[nb] is not NULL
         ORDER BY find_id, treatment_id;
     """
+    TABLE_COLS = ['treatment__treatment_type',
+                  'treatment__upstream',
+                  'treatment__start_date', 'treatment__end_date',
+                  'treatment__location', 'treatment__container',
+                  'treatment__person', 'treatment_nb']
 
     class Meta:
         managed = False
-        db_table = 'find_downtreatments'
+        db_table = 'find_uptreatments'
         unique_together = ('find', 'treatment')
+        ordering = ('find', '-treatment_nb')
 
 
-class FindUpstreamTreatments(AbsFindTreatments):
+class FindDownstreamTreatments(AbsFindTreatments):
     """
-    CREATE VIEW find_uptreatments_tree AS
+    CREATE VIEW find_downtreatments_tree AS
         WITH RECURSIVE rel_tree AS (
           SELECT id AS find_id, downstream_treatment_id, upstream_treatment_id,
               1 AS level,
@@ -834,19 +835,25 @@ class FindUpstreamTreatments(AbsFindTreatments):
         SELECT DISTINCT find_id, path_info, level
           FROM rel_tree ORDER BY find_id;
 
-    CREATE VIEW find_uptreatments AS
+    CREATE VIEW find_downtreatments AS
         SELECT DISTINCT find_id,
             path_info[nb] AS treatment_id, level - nb + 1 AS treatment_nb
           FROM (SELECT *, generate_subscripts(path_info, 1) AS nb
-                FROM find_uptreatments_tree) y
+                FROM find_downtreatments_tree) y
          WHERE path_info[nb] is not NULL
         ORDER BY find_id, treatment_id;
     """
+    TABLE_COLS = ['treatment__treatment_type',
+                  'treatment__downstream',
+                  'treatment__start_date', 'treatment__end_date',
+                  'treatment__location', 'treatment__container',
+                  'treatment__person', 'treatment_nb']
 
     class Meta:
         managed = False
-        db_table = 'find_uptreatments'
+        db_table = 'find_downtreatments'
         unique_together = ('find', 'treatment')
+        ordering = ('find', '-treatment_nb')
 
 
 class FindTreatments(AbsFindTreatments):
@@ -865,6 +872,7 @@ class FindTreatments(AbsFindTreatments):
         managed = False
         db_table = 'find_treatments'
         unique_together = ('find', 'treatment')
+        ordering = ('find', 'upstream', '-treatment_nb')
 
 
 class TreatmentSource(Source):