add geography field to dataset and usecase. Add sdg to usecase

Author: Deepthi-Chand

api/models/Dataset.py

@@ -59,6 +59,9 @@ class Dataset(models.Model):
         max_length=50, default=DatasetStatus.DRAFT, choices=DatasetStatus.choices
     )
     sectors = models.ManyToManyField("api.Sector", blank=True, related_name="datasets")
+    geographies = models.ManyToManyField(
+        "api.Geography", blank=True, related_name="datasets"
+    )
     access_type = models.CharField(
         max_length=50,
         default=DatasetAccessType.PUBLIC,
@@ -91,6 +94,14 @@ def sectors_indexing(self) -> list[str]:
         """
         return [sector.name for sector in self.sectors.all()]  # type: ignore
 
+    @property
+    def geographies_indexing(self) -> list[str]:
+        """Geographies for indexing.
+
+        Used in Elasticsearch indexing.
+        """
+        return [geo.name for geo in self.geographies.all()]  # type: ignore
+
     @property
     def formats_indexing(self) -> list[str]:
         """Formats for indexing.

api/models/UseCase.py

@@ -44,6 +44,10 @@ class UseCase(models.Model):
         choices=UseCaseRunningStatus.choices,
     )
     sectors = models.ManyToManyField("api.Sector", blank=True, related_name="usecases")
+    sdgs = models.ManyToManyField("api.SDG", blank=True, related_name="usecases")
+    geographies = models.ManyToManyField(
+        "api.Geography", blank=True, related_name="usecases"
+    )
     contributors = models.ManyToManyField(
         "authorization.User", blank=True, related_name="contributed_usecases"
     )
@@ -75,5 +79,13 @@ def sectors_indexing(self):
     def tags_indexing(self):
         return [tag.value for tag in self.tags.all()]  # type: ignore
 
+    @property
+    def sdgs_indexing(self):
+        return [sdg.code for sdg in self.sdgs.all()]  # type: ignore
+
+    @property
+    def geographies_indexing(self):
+        return [geo.name for geo in self.geographies.all()]  # type: ignore
+
     class Meta:
         db_table = "use_case"

api/schema/dataset_schema.py

@@ -12,6 +12,7 @@
 
 from api.models import (
     Dataset,
+    Geography,
     Metadata,
     Organization,
     Resource,
@@ -253,6 +254,7 @@ class UpdateMetadataInput:
     description: Optional[str] = None
     tags: Optional[List[str]] = None
     sectors: Optional[List[uuid.UUID]] = None
+    geographies: Optional[List[int]] = None
     access_type: Optional[DatasetAccessTypeENUM] = DatasetAccessTypeENUM.PUBLIC
     license: Optional[DatasetLicenseENUM] = (
         DatasetLicenseENUM.CC_BY_SA_4_0_ATTRIBUTION_SHARE_ALIKE
@@ -328,6 +330,17 @@ def _add_update_dataset_sectors(dataset: Dataset, sectors: List[uuid.UUID]) -> N
     dataset.save()
 
 
+@trace_resolver(
+    name="add_update_dataset_geographies", attributes={"component": "dataset"}
+)
+def _add_update_dataset_geographies(dataset: Dataset, geography_ids: List[int]) -> None:
+    """Update geographies for a dataset."""
+    dataset.geographies.clear()
+    geographies = Geography.objects.filter(id__in=geography_ids)
+    dataset.geographies.add(*geographies)
+    dataset.save()
+
+
 @strawberry.type
 class Query:
     @strawberry_django.field(
@@ -595,6 +608,8 @@ def add_update_dataset_metadata(
         _add_update_dataset_metadata(dataset, metadata_input)
         if update_metadata_input.sectors is not None:
             _add_update_dataset_sectors(dataset, update_metadata_input.sectors)
+        if update_metadata_input.geographies is not None:
+            _add_update_dataset_geographies(dataset, update_metadata_input.geographies)
         return MutationResponse.success_response(TypeDataset.from_django(dataset))
 
     @strawberry_django.mutation(

api/schema/usecase_schema.py

@@ -16,7 +16,9 @@
 from strawberry_django.pagination import OffsetPaginationInput
 
 from api.models import (
+    SDG,
     Dataset,
+    Geography,
     Metadata,
     Organization,
     Sector,
@@ -62,6 +64,8 @@ class UpdateUseCaseMetadataInput:
     metadata: List[UCMetadataItemType]
     tags: Optional[List[str]]
     sectors: List[uuid.UUID]
+    sdgs: Optional[List[uuid.UUID]]
+    geographies: Optional[List[int]] = None
 
 
 use_case_running_status = strawberry.enum(UseCaseRunningStatus)  # type: ignore
@@ -79,6 +83,8 @@ class UseCaseInputPartial:
     platform_url: Optional[str] = None
     tags: Optional[List[str]] = None
     sectors: Optional[List[uuid.UUID]] = None
+    sdgs: Optional[List[uuid.UUID]] = None
+    geographies: Optional[List[int]] = None
     started_on: Optional[datetime.date] = None
     completed_on: Optional[datetime.date] = None
 
@@ -202,6 +208,23 @@ def _update_usecase_sectors(usecase: UseCase, sectors: List[uuid.UUID]) -> None:
     usecase.save()
 
 
+@trace_resolver(name="update_usecase_geographies", attributes={"component": "usecase"})
+def _update_usecase_geographies(usecase: UseCase, geography_ids: List[int]) -> None:
+    """Update geographies for a usecase."""
+    usecase.geographies.clear()
+    geographies = Geography.objects.filter(id__in=geography_ids)
+    usecase.geographies.add(*geographies)
+    usecase.save()
+
+
+@trace_resolver(name="update_usecase_sdgs", attributes={"component": "usecase"})
+def _update_usecase_sdgs(usecase: UseCase, sdgs: List[uuid.UUID]) -> None:
+    sdgs_objs = SDG.objects.filter(id__in=sdgs)
+    usecase.sdgs.clear()
+    usecase.sdgs.add(*sdgs_objs)
+    usecase.save()
+
+
 @trace_resolver(
     name="add_update_usecase_metadata",
     attributes={"component": "usecase", "operation": "mutation"},
@@ -339,6 +362,10 @@ def add_update_usecase_metadata(
             _update_usecase_tags(usecase, update_metadata_input.tags)
         _add_update_usecase_metadata(usecase, metadata_input)
         _update_usecase_sectors(usecase, update_metadata_input.sectors)
+        if update_metadata_input.sdgs is not None:
+            _update_usecase_sdgs(usecase, update_metadata_input.sdgs)
+        if update_metadata_input.geographies is not None:
+            _update_usecase_geographies(usecase, update_metadata_input.geographies)
         return TypeUseCase.from_django(usecase)
 
     @strawberry_django.mutation(handle_django_errors=False)

api/types/type_dataset.py

@@ -11,6 +11,7 @@
 from api.models import Dataset, DatasetMetadata, Resource, Tag
 from api.types.base_type import BaseType
 from api.types.type_dataset_metadata import TypeDatasetMetadata
+from api.types.type_geo import TypeGeo
 from api.types.type_organization import TypeOrganization
 from api.types.type_resource import TypeResource
 from api.types.type_sector import TypeSector
@@ -78,6 +79,25 @@ def sectors(self, info: Info) -> List["TypeSector"]:
         except (AttributeError, Dataset.DoesNotExist):
             return []
 
+    @strawberry.field
+    def geographies(self, info: Info) -> List["TypeGeo"]:
+        """Get geographies for this dataset.
+
+        Args:
+            info: Request info
+
+        Returns:
+            List[TypeGeo]: List of geographies
+        """
+        try:
+            from api.types.type_geo import TypeGeo
+
+            django_instance = cast(Dataset, self)
+            queryset = django_instance.geographies.all()
+            return TypeGeo.from_django_list(queryset)
+        except (AttributeError, Dataset.DoesNotExist):
+            return []
+
     @strawberry.field
     def metadata(self) -> List["TypeDatasetMetadata"]:
         """Get metadata for this dataset."""

api/types/type_usecase.py

@@ -15,7 +15,9 @@
 )
 from api.types.base_type import BaseType
 from api.types.type_dataset import TypeDataset, TypeTag
+from api.types.type_geo import TypeGeo
 from api.types.type_organization import TypeOrganization
+from api.types.type_sdg import TypeSDG
 from api.types.type_sector import TypeSector
 from api.types.type_usecase_dashboard import TypeUseCaseDashboard
 from api.types.type_usecase_metadata import TypeUseCaseMetadata
@@ -113,6 +115,28 @@ def sectors(self) -> Optional[List["TypeSector"]]:
         except Exception:
             return []
 
+    @strawberry.field(description="Get SDGs associated with this use case.")
+    def sdgs(self) -> Optional[List["TypeSDG"]]:
+        """Get SDGs associated with this use case."""
+        try:
+            queryset = self.sdgs.all()  # type: ignore
+            if not queryset.exists():
+                return []
+            return TypeSDG.from_django_list(queryset)
+        except Exception:
+            return []
+
+    @strawberry.field(description="Get geographies associated with this use case.")
+    def geographies(self) -> Optional[List["TypeGeo"]]:
+        """Get geographies associated with this use case."""
+        try:
+            queryset = self.geographies.all()  # type: ignore
+            if not queryset.exists():
+                return []
+            return TypeGeo.from_django_list(queryset)
+        except Exception:
+            return []
+
     @strawberry.field(description="Get tags associated with this use case.")
     def tags(self) -> Optional[List["TypeTag"]]:
         """Get tags associated with this use case."""

api/views/search_dataset.py

@@ -64,6 +64,7 @@ class DatasetDocumentSerializer(serializers.ModelSerializer):
     sectors = serializers.ListField()
     formats = serializers.ListField()
     catalogs = serializers.ListField()
+    geographies = serializers.ListField()
     has_charts = serializers.BooleanField()
     slug = serializers.CharField()
     download_count = serializers.IntegerField()
@@ -97,6 +98,7 @@ class Meta:
             "sectors",
             "formats",
             "catalogs",
+            "geographies",
             "has_charts",
             "download_count",
             "trending_score",
@@ -145,6 +147,7 @@ def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
             "sectors.raw": "terms",
             "formats.raw": "terms",
             "catalogs.raw": "terms",
+            "geographies.raw": "terms",
         }
         for metadata in enabled_metadata:  # type: Metadata
             if metadata.filterable:
@@ -253,7 +256,7 @@ def add_filters(self, filters: Dict[str, str], search: Search) -> Search:
         for filter in filters:
             if filter in excluded_labels:
                 continue
-            elif filter in ["tags", "sectors", "formats", "catalogs"]:
+            elif filter in ["tags", "sectors", "formats", "catalogs", "geographies"]:
                 raw_filter = filter + ".raw"
                 if raw_filter in self.aggregations:
                     filter_values = filters[filter].split(",")

api/views/search_usecase.py

@@ -63,6 +63,7 @@ class UseCaseDocumentSerializer(serializers.ModelSerializer):
     tags = serializers.ListField()
     logo = serializers.CharField()
     sectors = serializers.ListField(default=[])
+    geographies = serializers.ListField(default=[])
     slug = serializers.CharField()
     is_individual_usecase = serializers.BooleanField()
     running_status = serializers.CharField()
@@ -117,6 +118,7 @@ class Meta:
             "tags",
             "logo",
             "sectors",
+            "geographies",
             "is_individual_usecase",
             "organization",
             "user",
@@ -169,6 +171,7 @@ def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
         aggregations: Dict[str, str] = {
             "tags.raw": "terms",
             "sectors.raw": "terms",
+            "geographies.raw": "terms",
             "status": "terms",
             "running_status": "terms",
         }
@@ -304,7 +307,7 @@ def add_filters(self, filters: Dict[str, str], search: Search) -> Search:
         for filter in filters:
             if filter in excluded_labels:
                 continue
-            elif filter in ["tags", "sectors"]:
+            elif filter in ["tags", "sectors", "geographies"]:
                 raw_filter = filter + ".raw"
                 if raw_filter in self.aggregations:
                     filter_values = filters[filter].split(",")

search/documents/dataset_document.py

@@ -6,6 +6,7 @@
     Catalog,
     Dataset,
     DatasetMetadata,
+    Geography,
     Metadata,
     Organization,
     Resource,
@@ -114,6 +115,16 @@ class DatasetDocument(Document):
         multi=True,
     )
 
+    geographies = fields.TextField(
+        attr="geographies_indexing",
+        analyzer=ngram_analyser,
+        fields={
+            "raw": fields.KeywordField(multi=True),
+            "suggest": fields.CompletionField(multi=True),
+        },
+        multi=True,
+    )
+
     is_individual_dataset = fields.BooleanField(attr="is_individual_dataset")
 
     has_charts = fields.BooleanField(attr="has_charts")
@@ -185,7 +196,14 @@ def get_queryset(self) -> Any:
     def get_instances_from_related(
         self,
         related_instance: Union[
-            Resource, Metadata, DatasetMetadata, Sector, Organization, User, Catalog
+            Resource,
+            Metadata,
+            DatasetMetadata,
+            Sector,
+            Organization,
+            User,
+            Catalog,
+            Geography,
         ],
     ) -> Optional[Union[Dataset, List[Dataset]]]:
         """Get Dataset instances from related models."""
@@ -204,6 +222,8 @@ def get_instances_from_related(
             return list(related_instance.datasets.all())
         elif isinstance(related_instance, Catalog):
             return list(related_instance.datasets.all())
+        elif isinstance(related_instance, Geography):
+            return list(related_instance.datasets.all())
         return None
 
     class Django:
@@ -225,4 +245,5 @@ class Django:
             Organization,
             User,
             Catalog,
+            Geography,
         ]