Merge pull request #804 from basedosdados/feat/search-ajust

feat(search): add new filter in Elastic Search

Author: AldemirLucas

backend/apps/api/v1/admin.py

@@ -615,6 +615,10 @@ class DatasetAdmin(OrderedInlineModelAdminMixin, TabbedTranslationAdmin):
         "contains_raw_data_sources",
         "contains_information_requests",
         "contains_closed_data",
+        "contains_direct_download_free",
+        "contains_direct_download_paid",
+        "contains_temporalcoverage_free",
+        "contains_temporalcoverage_paid",
         "page_views",
         "created_at",
         "updated_at",
@@ -697,6 +701,10 @@ class TableAdmin(OrderedInlineModelAdminMixin, TabbedTranslationAdmin):
         "uncompressed_file_size",
         "compressed_file_size",
         "contains_open_data",
+        "contains_direct_download_free",
+        "contains_direct_download_paid",
+        "contains_temporalcoverage_free",
+        "contains_temporalcoverage_paid",
         "contains_closed_data",
         "page_views",
     ]

backend/apps/api/v1/models.py

@@ -680,6 +680,54 @@ def contains_closed_data(self):
                 return True
         return False
 
+    @property
+    def contains_direct_download_free(self):
+        return len(
+            [
+                table
+                for table in self.tables.exclude(status__slug__in=["under_review", "excluded"])
+                .exclude(slug__in=["dicionario", "dictionary"])
+                .all()
+                if table.contains_direct_download_free
+            ]
+        )
+
+    @property
+    def contains_direct_download_paid(self):
+        return len(
+            [
+                table
+                for table in self.tables.exclude(status__slug__in=["under_review", "excluded"])
+                .exclude(slug__in=["dicionario", "dictionary"])
+                .all()
+                if table.contains_direct_download_paid
+            ]
+        )
+
+    @property
+    def contains_temporalcoverage_free(self):
+        return len(
+            [
+                table
+                for table in self.tables.exclude(status__slug__in=["under_review", "excluded"])
+                .exclude(slug__in=["dicionario", "dictionary"])
+                .all()
+                if table.contains_temporalcoverage_free
+            ]
+        )
+
+    @property
+    def contains_temporalcoverage_paid(self):
+        return len(
+            [
+                table
+                for table in self.tables.exclude(status__slug__in=["under_review", "excluded"])
+                .exclude(slug__in=["dicionario", "dictionary"])
+                .all()
+                if table.contains_temporalcoverage_paid
+            ]
+        )
+
     @property
     def contains_tables(self):
         """Returns true if there are tables in the dataset"""
@@ -1098,6 +1146,28 @@ def contains_closed_data(self):
             return True
         return False
 
+    @property
+    def contains_direct_download_free(self):
+        if self.uncompressed_file_size is None:
+            return False
+        return self.uncompressed_file_size < 100 * 1024 * 1024
+
+    @property
+    def contains_direct_download_paid(self):
+        if self.uncompressed_file_size is None:
+            return False
+        return self.uncompressed_file_size > 100 * 1024 * 1024
+
+    @property
+    def contains_temporalcoverage_free(self) -> bool:
+        coverage = get_full_temporal_coverage([self]) or []
+        return all(entry["type"] == "open" for entry in coverage)
+
+    @property
+    def contains_temporalcoverage_paid(self) -> bool:
+        coverage = get_full_temporal_coverage([self]) or []
+        return any(entry["type"] == "closed" for entry in coverage)
+
     @property
     def temporal_coverage(self) -> dict:
         """Temporal coverage"""

backend/apps/api/v1/schemas.py

@@ -81,6 +81,11 @@ class Dataset(BaseModel):
     contains_open_data: bool
     contains_closed_data: bool
     #
+    contains_direct_download_free: bool
+    contains_direct_download_paid: bool
+    contains_temporalcoverage_free: bool
+    contains_temporalcoverage_paid: bool
+    #
     themes: List[Theme]
     organization: List[Organization]
     temporal_coverage: List[str]

backend/apps/api/v1/search_indexes.py

@@ -242,6 +242,26 @@ class DatasetIndex(indexes.SearchIndex, indexes.Indexable):
         indexed=False,
     )
 
+    contains_direct_download_free = indexes.BooleanField(
+        model_attr="contains_direct_download_free",
+        indexed=False,
+    )
+
+    contains_direct_download_paid = indexes.BooleanField(
+        model_attr="contains_direct_download_paid",
+        indexed=False,
+    )
+
+    contains_temporalcoverage_free = indexes.BooleanField(
+        model_attr="contains_temporalcoverage_free",
+        indexed=False,
+    )
+
+    contains_temporalcoverage_paid = indexes.BooleanField(
+        model_attr="contains_temporalcoverage_paid",
+        indexed=False,
+    )
+
     contains_tables = indexes.BooleanField(
         model_attr="contains_tables",
         indexed=False,

backend/apps/api/v1/search_views.py

@@ -2,6 +2,7 @@
 
 from django.core.files.storage import default_storage as storage
 from django.http import JsonResponse
+from django.views import View
 from haystack.forms import FacetedSearchForm
 from haystack.generic_views import FacetedSearchView
 from haystack.models import SearchResult
@@ -30,17 +31,6 @@ def search(self):
         # Start with all results
         sqs = self.searchqueryset.all()
 
-        # Debug print to see all form data
-        print(
-            "DEBUG: Form data:",
-            {
-                "spatial_coverage": self.spatial_coverage,
-                "theme": self.theme,
-                "organization": self.organization,
-                "tag": self.tag,
-            },
-        )
-
         # Text search if provided
         if q := self.cleaned_data.get("q"):
             sqs = (
@@ -109,6 +99,10 @@ class DatasetSearchView(FacetedSearchView):
         "contains_tables",
         "contains_raw_data_sources",
         "contains_information_requests",
+        "contains_direct_download_free",
+        "contains_direct_download_paid",
+        "contains_temporalcoverage_free",
+        "contains_temporalcoverage_paid",
     ]
 
     @property
@@ -153,10 +147,10 @@ def get(self, request, *args, **kwargs):
             }
         )
 
-    def get_facets(self, sqs: SearchQuerySet, facet_size=200):
+    def get_facets(self, sqs: SearchQuerySet, facet_size=6):
         sqs = sqs.facet("theme_slug", size=facet_size)
         sqs = sqs.facet("organization_slug", size=facet_size)
-        sqs = sqs.facet("spatial_coverage", size=facet_size)
+        # sqs = sqs.facet("spatial_coverage", size=facet_size)
         sqs = sqs.facet("tag_slug", size=facet_size)
         sqs = sqs.facet("entity_slug", size=facet_size)
 
@@ -296,9 +290,7 @@ def as_search_result(result: SearchResult, locale="pt"):
         )
 
     entities = []
-    for slug, name in zip(
-        result.entity_slug or [], getattr(result, f"entity_name_{locale}") or []
-    ):
+    for slug, name in zip(result.entity_slug or [], getattr(result, f"entity_name_{locale}") or []):
         entities.append(
             {
                 "slug": slug,
@@ -338,6 +330,10 @@ def as_search_result(result: SearchResult, locale="pt"):
         "contains_open_data": result.contains_open_data,
         "contains_closed_data": result.contains_closed_data,
         "contains_tables": result.contains_tables,
+        "contains_direct_download_free": result.contains_direct_download_free,
+        "contains_direct_download_paid": result.contains_direct_download_paid,
+        "contains_temporalcoverage_free": result.contains_temporalcoverage_free,
+        "contains_temporalcoverage_paid": result.contains_temporalcoverage_paid,
         "contains_raw_data_sources": result.contains_raw_data_sources,
         "contains_information_requests": result.contains_information_requests,
         "n_tables": result.n_tables,
@@ -349,3 +345,104 @@ def as_search_result(result: SearchResult, locale="pt"):
         "first_raw_data_source_id": result.first_raw_data_source_id,
         "first_information_request_id": result.first_information_request_id,
     }
+
+
+class DatasetFacetValuesView(View):
+    """
+    View para retornar os valores de uma faceta específica baseada nos parâmetros de filtro atuais.
+    """
+
+    facet_fields = [
+        "tag_slug",
+        "theme_slug",
+        "entity_slug",
+        "organization_slug",
+        "spatial_coverage",
+    ]
+
+    @property
+    def locale(self):
+        return self.request.GET.get("locale", "pt")
+
+    @property
+    def facet_name(self):
+        return self.request.GET.get("facet", "").lower()
+
+    def get_form_kwargs(self):
+        kwargs = {
+            "contains": self.request.GET.getlist("contains"),
+            "theme": self.request.GET.getlist("theme"),
+            "organization": self.request.GET.getlist("organization"),
+            "spatial_coverage": self.request.GET.getlist("spatial_coverage"),
+            "tag": self.request.GET.getlist("tag"),
+            "observation_level": self.request.GET.getlist("observation_level"),
+            "locale": self.locale,
+        }
+
+        if "q" in self.request.GET:
+            kwargs["data"] = {"q": self.request.GET["q"]}
+
+        return kwargs
+
+    def get(self, request, *args, **kwargs):
+        if self.facet_name not in self.facet_fields:
+            return JsonResponse(
+                {
+                    "error": f"Facet '{self.facet_name}' not available. Choose from: {self.facet_fields}"
+                },
+                status=400,
+            )
+
+        form = DatasetSearchForm(**self.get_form_kwargs())
+        sqs = form.search()
+
+        sqs = sqs.facet(self.facet_name, size=1000)  # Tamanho grande para pegar todos os valores
+
+        facet_counts = sqs.facet_counts()
+        facet_values = facet_counts.get("fields", {}).get(self.facet_name, [])
+
+        values = [{"key": value[0], "count": value[1]} for value in facet_values if value[0]]
+
+        if self.facet_name == "theme_slug":
+            model = Theme
+            name_field = f"name_{self.locale}"
+        elif self.facet_name == "organization_slug":
+            model = Organization
+            name_field = f"name_{self.locale}"
+        elif self.facet_name == "tag_slug":
+            model = Tag
+            name_field = f"name_{self.locale}"
+        elif self.facet_name == "entity_slug":
+            model = Entity
+            name_field = f"name_{self.locale}"
+        elif self.facet_name == "spatial_coverage":
+            model = Area
+            name_field = f"name_{self.locale}"
+        else:
+            model = None
+
+        if model:
+            slugs = [v["key"] for v in values]
+            translated_names = model.objects.filter(slug__in=slugs).values(
+                "slug", name_field, "name"
+            )
+            name_map = {
+                item["slug"]: {
+                    "name": item.get(name_field) or item.get("name") or item["slug"],
+                    "fallback": name_field not in item or item[name_field] is None,
+                }
+                for item in translated_names
+            }
+
+            for value in values:
+                translated = name_map.get(value["key"], {"name": value["key"], "fallback": True})
+                value.update(translated)
+
+        return JsonResponse(
+            {
+                "facet": self.facet_name,
+                "values": values,
+                "count": len(values),
+                "locale": self.locale,
+            }
+        )

backend/apps/api/v1/urls.py

@@ -4,7 +4,7 @@
 from django.views.decorators.csrf import csrf_exempt
 from graphene_file_upload.django import FileUploadGraphQLView
 
-from backend.apps.api.v1.search_views import DatasetSearchView
+from backend.apps.api.v1.search_views import DatasetFacetValuesView, DatasetSearchView
 from backend.apps.api.v1.views import DatasetRedirectView
 
 
@@ -22,6 +22,7 @@ def graphql_view():
     path("api/v1/graphql", graphql_view()),
     path("graphql", graphql_view()),
     path("search/", DatasetSearchView.as_view()),
+    path("facet_values/", DatasetFacetValuesView.as_view()),
     path("dataset/", DatasetRedirectView.as_view()),
     path("dataset_redirect/", DatasetRedirectView.as_view()),
 ]