From 962cf0a79e26e44c88fe9e1e5b2ffcc5dec76ed4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Formento?= <andreformento.sc@gmail.com>
Date: Tue, 20 Apr 2021 22:43:48 -0300
Subject: [PATCH] create autocomplete endpoint

---
 api/api.py                 | 36 ++++++++++++++++++++++++++++++
 database/elasticsearch.py  | 33 +++++++++++++++++++++++++++-
 gazettes/__init__.py       |  1 +
 gazettes/gazette_access.py | 45 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 114 insertions(+), 1 deletion(-)

diff --git a/api/api.py b/api/api.py
index 2b0fab6..7feaa1b 100644
--- a/api/api.py
+++ b/api/api.py
@@ -28,6 +28,16 @@ class GazetteSearchResponse(BaseModel):
     gazettes: List[GazetteItem]
 
 
+class CityAutocompleteItem(BaseModel):
+    territory_id: str
+    territory_name: str
+    state_code: str
+
+
+class CityAutocompleteResponse(BaseModel):
+    suggestions: List[CityAutocompleteItem]
+
+
 def trigger_gazettes_search(
     territory_id: str = None,
     since: date = None,
@@ -128,6 +138,32 @@ async def get_gazettes_by_territory_id(
 ):
     return trigger_gazettes_search(territory_id, since, until, keywords, offset, size)
 
+@app.get(
+    "/cities/autocomplete",
+    response_model=CityAutocompleteResponse,
+    name="Autocomplete city",
+    description="Autocomplete city by term",
+    response_model_exclude_unset=True,
+    response_model_exclude_none=True,
+)
+async def autocomplete_city(
+    term: str = Query(
+        0, title="Term", description="Term to autocomplete",
+    ),
+    size: Optional[int] = Query(
+        10,
+        title="Number of item to return",
+        description="Define the number of item should be returned",
+    ),
+):
+    def convert(item):
+        return CityAutocompleteItem(
+            territory_id=item.territory_id,
+            territory_name=item.territory_name,
+            state_code=item.state_code)
+
+    suggestions = list(map(convert, app.gazettes.autocomplete_city(term, size)))
+    return CityAutocompleteResponse(suggestions=suggestions)
 
 def configure_api_app(gazettes: GazetteAccessInterface, api_root_path=None):
     if not isinstance(gazettes, GazetteAccessInterface):
diff --git a/database/elasticsearch.py b/database/elasticsearch.py
index f784267..b4618d6 100644
--- a/database/elasticsearch.py
+++ b/database/elasticsearch.py
@@ -4,7 +4,7 @@
 
 import elasticsearch
 
-from gazettes import GazetteDataGateway, Gazette
+from gazettes import GazetteDataGateway, Gazette, CityAutocomplete
 
 
 class ElasticSearchDataMapper(GazetteDataGateway):
@@ -120,6 +120,37 @@ def get_gazettes(
             self.create_list_with_gazette_objects(gazettes["hits"]["hits"]),
         )
 
+    def autocomplete_city(
+        self,
+        term=None,
+        size=10,
+    ):
+        query = { "query": { "match_phrase_prefix": { "territory_name": { "query": term } } },
+                  "size": 0,
+                  "aggs": {
+                    "territory_and_state": {
+                        "composite": {
+                        "size": size,
+                        "sources": [
+                            { "territory_id": { "terms": { "field": "territory_id.keyword" } } },
+                            { "territory_name": { "terms": { "field": "territory_name.keyword" } } },
+                            { "state_code": { "terms": { "field": "state_code.keyword" } } }
+                            ]
+                        }
+                    }
+                }
+               }
+        search_result = self._es.search(body=query, index=self._index)
+
+        def convert(document):
+            return CityAutocomplete(
+                document["key"]["territory_id"],
+                document["key"]["territory_name"],
+                document["key"]["state_code"],
+            )
+
+        return list(map(convert, search_result["aggregations"]["territory_and_state"]["buckets"]))
+
 
 def create_elasticsearch_data_mapper(
     host: str = None, index: str = None
diff --git a/gazettes/__init__.py b/gazettes/__init__.py
index dd69b49..8bcb98e 100644
--- a/gazettes/__init__.py
+++ b/gazettes/__init__.py
@@ -4,5 +4,6 @@
     GazetteRequest,
     GazetteDataGateway,
     Gazette,
+    CityAutocomplete,
     create_gazettes_interface,
 )
diff --git a/gazettes/gazette_access.py b/gazettes/gazette_access.py
index 4cd05c3..5b96799 100644
--- a/gazettes/gazette_access.py
+++ b/gazettes/gazette_access.py
@@ -38,6 +38,14 @@ def get_gazettes(
         Method to get the gazette from storage
         """
 
+    @abc.abstractmethod
+    def autocomplete_city(
+        self, term=None, size: int = 10
+    ):
+        """
+        Method to autocomplete city name
+        """
+
 
 class GazetteAccessInterface(abc.ABC):
     """
@@ -75,6 +83,9 @@ def get_gazettes(self, filters: GazetteRequest = None):
         )
         return (total_number_gazettes, [vars(gazette) for gazette in gazettes])
 
+    def autocomplete_city(self, term=None, size: int = 10):
+        return self._data_gateway.autocomplete_city(term, size)
+
 
 class Gazette:
     """
@@ -130,6 +141,40 @@ def __eq__(self, other):
     def __repr__(self):
         return f"Gazette({self.checksum}, {self.territory_id}, {self.date}, {self.url}, {self.territory_name}, {self.state_code}, {self.edition}, {self.is_extra_edition})"
 
+class CityAutocomplete:
+    """
+    Item to represent a city
+    """
+
+    def __init__(
+        self,
+        territory_id,
+        territory_name,
+        state_code,
+    ):
+        self.territory_id = territory_id
+        self.territory_name = territory_name
+        self.state_code = state_code
+
+    def __hash__(self):
+        return hash(
+            (
+                self.territory_id,
+                self.territory_name,
+                self.state_code,
+            )
+        )
+
+    def __eq__(self, other):
+        return (
+            self.territory_id == other.territory_id
+            and self.territory_name == other.territory_name
+            and self.state_code == other.state_code
+        )
+
+    def __repr__(self):
+        return f"CityAutocomplete({self.territory_id}, {self.territory_name}, {self.state_code})"
+
 
 def create_gazettes_interface(data_gateway: GazetteDataGateway):
     if not isinstance(data_gateway, GazetteDataGateway):