From e3cbbe735685d3e8834abd1c676bfd582eb18171 Mon Sep 17 00:00:00 2001
From: Struan Donald <struan@exo.org.uk>
Date: Thu, 9 Nov 2023 11:45:27 +0000
Subject: [PATCH] update dataset areas_avaible in base importers

---
 hub/management/commands/base_importers.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/hub/management/commands/base_importers.py b/hub/management/commands/base_importers.py
index 8eadc564b..95036a621 100644
--- a/hub/management/commands/base_importers.py
+++ b/hub/management/commands/base_importers.py
@@ -1,3 +1,4 @@
+from functools import cache
 from time import sleep
 
 from django.core.management.base import BaseCommand
@@ -7,7 +8,7 @@
 import pandas as pd
 from tqdm import tqdm
 
-from hub.models import Area, AreaData, DataSet, DataType
+from hub.models import Area, AreaData, AreaType, DataSet, DataType
 from utils.mapit import (
     BadRequestException,
     ForbiddenException,
@@ -20,6 +21,7 @@
 
 class BaseAreaImportCommand(BaseCommand):
     cast_field = IntegerField
+    area_type = "WMC"
 
     def __init__(self):
         super().__init__()
@@ -37,6 +39,10 @@ def get_label(self, config):
     def delete_data(self):
         pass
 
+    @cache
+    def get_area_type(self):
+        return AreaType.objects.get(code=self.area_type)
+
     def add_data_sets(self, df=None):
         for name, config in self.data_sets.items():
             label = self.get_label(config)
@@ -69,10 +75,12 @@ def add_data_sets(self, df=None):
                     **config["defaults"],
                 },
             )
+            data_set.areas_available.add(self.get_area_type())
 
             data_type, created = DataType.objects.update_or_create(
                 data_set=data_set,
                 name=name,
+                area_type=self.get_area_type(),
                 defaults={
                     "data_type": config["defaults"]["data_type"],
                     "label": label,
@@ -163,7 +171,6 @@ def update_max_min(self):
 
 class BaseImportFromDataFrameCommand(BaseAreaImportCommand):
     uses_gss = True
-    area_type = "WMC"
 
     def get_row_data(self, row, conf):
         return row[conf["col"]]