Fix US harvester and some commands (#122)

harvesters/harvester/base.py

@@ -170,7 +170,8 @@ def _save_well(
             feature_type: typing.Optional[TermFeatureType] = None,
             ground_surface_elevation_masl: typing.Optional[float] = None,
             top_of_well_elevation_masl: typing.Optional[float] = None,
-            description: typing.Optional[str] = None
+            description: typing.Optional[str] = None,
+            reassign_organisation=False
     ):
         """ Save well """
         well = self.get_well(
@@ -233,6 +234,10 @@ def _save_well(
             harvester=self.harvester,
             well=well
         )
+        if reassign_organisation:
+            if well.organisation != self.harvester.organisation:
+                well.organisation = self.harvester.organisation
+                well.save()
         return well, harvester_well_data
 
     def _save_measurement(

harvesters/harvester/cida.py

@@ -42,6 +42,8 @@ def __init__(
             self.units['feet'] = Unit.objects.get(name='ft')
             self.units['ft'] = Unit.objects.get(name='ft')
             self.units['m'] = Unit.objects.get(name='m')
+            self.units['meters'] = Unit.objects.get(name='m')
+            self.units['centimeters'] = Unit.objects.get(name='cm')
         except Unit.DoesNotExist:
             pass
 
@@ -154,7 +156,8 @@ def get_stations(self, xml):
                     name=site_name,
                     latitude=latitude,
                     longitude=longitude,
-                    ground_surface_elevation_masl=altitude
+                    ground_surface_elevation_masl=altitude,
+                    reassign_organisation=True
                 )
 
                 # Total depth

management/commands/generate_data_wells_cache.py

@@ -61,8 +61,10 @@ def handle(self, *args, **options):
         countries = []
         organisations = []
         count = wells.count()
-        for idx, well in enumerate(wells.order_by('id')):
-            print(f'----- {idx}/{count} -----')
+        ids = list(wells.order_by('id').values_list('id', flat=True))
+        for idx, id in enumerate(ids):
+            well = Well.objects.get(id=id)
+            print(f'----- {idx}/{count} - {well.id} -----')
             generate_data_well_cache(
                 well.id, force_regenerate=force, generate_country_cache=False,
                 generate_organisation_cache=False

management/commands/generate_well_measurement_cache.py

@@ -1,5 +1,7 @@
 import os
+
 from django.core.management.base import BaseCommand
+
 from gwml2.models.well import (
     Well, WellLevelMeasurement, WellYieldMeasurement, WellQualityMeasurement)
 
@@ -29,20 +31,46 @@ def add_arguments(self, parser):
             default='',
             dest='force',
             help='Force to regenerate')
+        parser.add_argument(
+            '-from_id',
+            '--from_id',
+            dest='from_id',
+            help='From id'
+        )
+        parser.add_argument(
+            '-country_code',
+            '--country_code',
+            dest='country_code',
+            help='From country code'
+        )
 
     def handle(self, *args, **options):
         id = options.get('id', None)
         measurement_name = options.get('measurement_name', None)
+        from_id = options.get('from_id', False)
+        country_code = options.get('country_code', False)
         force = options.get('force', False)
+
+        # Filter by from_id
         if id:
-            queryset = Well.objects.filter(id=int(id))
+            wells = Well.objects.filter(id=id)
+        elif from_id:
+            wells = Well.objects.filter(id__gte=from_id)
         else:
-            queryset = Well.objects.all()
+            wells = Well.objects.all()
+
+        # Check country code
+        if country_code:
+            wells = wells.filter(country__code=country_code)
 
-        count = queryset.count()
-        for idx, well in enumerate(queryset):
-            for MeasurementModel in \
-                    [WellLevelMeasurement, WellQualityMeasurement, WellYieldMeasurement]:
+        count = wells.count()
+        ids = list(wells.order_by('id').values_list('id', flat=True))
+        for idx, id in enumerate(ids):
+            well = Well.objects.get(id=id)
+            for MeasurementModel in [
+                WellLevelMeasurement, WellQualityMeasurement,
+                WellYieldMeasurement
+            ]:
                 # skip if measurement filtered
                 if measurement_name and MeasurementModel.__name__ != measurement_name:
                     continue
@@ -52,5 +80,7 @@ def handle(self, *args, **options):
                     if os.path.exists(cache_file):
                         continue
 
-                print("{}/{} : Generating {}, well {}".format(idx + 1, count, model, well.id))
+                print("{}/{} : Generating {}, well {}".format(
+                    idx + 1, count, model, well.id)
+                )
                 well.generate_measurement_cache(model)