Merge pull request #1294 from syucream/fix/attr-hint-type

Refine AttrHint type

Author: userlocalhost

airone/lib/elasticsearch.py

@@ -51,12 +51,12 @@ class FilterKey(enum.IntEnum):
     DUPLICATED = 5
 
 
-class AttrHint(TypedDict):
+class AttrHint(BaseModel):
     name: str
-    is_readable: NotRequired[bool]
-    filter_key: NotRequired[FilterKey]
-    keyword: NotRequired[str]
-    exact_match: NotRequired[bool]
+    is_readable: bool | None = None
+    filter_key: FilterKey | None = None
+    keyword: str | None = None
+    exact_match: bool | None = None
 
 
 class AttributeDocument(TypedDict):
@@ -267,23 +267,23 @@ def make_query(
 
     # Conversion processing from "filter_key" to "keyword" for each hint_attrs
     for hint_attr in hint_attrs:
-        match hint_attr.get("filter_key", None):
+        match hint_attr.filter_key:
             case FilterKey.CLEARED:
                 # remove "keyword" parameter
-                hint_attr.pop("keyword", None)
+                hint_attr.keyword = None
             case FilterKey.EMPTY:
-                hint_attr["keyword"] = "\\"
+                hint_attr.keyword = "\\"
             case FilterKey.NON_EMPTY:
-                hint_attr["keyword"] = "*"
+                hint_attr.keyword = "*"
             case FilterKey.DUPLICATED:
-                aggs_query = _make_aggs_query(hint_attr["name"])
+                aggs_query = _make_aggs_query(hint_attr.name)
                 # TODO Set to 1 for convenience
                 resp = execute_query(aggs_query, 1)
                 keyword_infos = resp["aggregations"]["attr_aggs"]["attr_name_aggs"][
                     "attr_value_aggs"
                 ]["buckets"]
                 keyword_list = [x["key"] for x in keyword_infos]
-                hint_attr["keyword"] = CONFIG.OR_SEARCH_CHARACTER.join(
+                hint_attr.keyword = CONFIG.OR_SEARCH_CHARACTER.join(
                     ["^" + x + "$" for x in keyword_list]
                 )
 
@@ -324,9 +324,7 @@ def make_query(
                     "query": {
                         "bool": {
                             "should": [
-                                {"term": {"attr.name": x["name"]}}
-                                for x in hint_attrs
-                                if "name" in x
+                                {"term": {"attr.name": x.name}} for x in hint_attrs if x.name
                             ]
                         }
                     },
@@ -337,7 +335,7 @@ def make_query(
     attr_query: dict[str, dict] = {}
 
     # filter attribute by keywords
-    for hint in [hint for hint in hint_attrs if "name" in hint and hint.get("keyword")]:
+    for hint in [hint for hint in hint_attrs if hint.name and hint.keyword]:
         attr_query.update(_parse_or_search(hint))
 
     # Build queries along keywords
@@ -649,7 +647,7 @@ def _parse_or_search(hint: AttrHint) -> dict[str, dict]:
     duplicate_keys: list = []
 
     # Split and process keywords with 'or'
-    for keyword_divided_or in hint["keyword"].split(CONFIG.OR_SEARCH_CHARACTER):
+    for keyword_divided_or in (hint.keyword or "").split(CONFIG.OR_SEARCH_CHARACTER):
         parsed_query = _parse_and_search(hint, keyword_divided_or, duplicate_keys)
         attr_query.update(parsed_query)
 
@@ -696,7 +694,7 @@ def _parse_and_search(
 
     # Keyword divided by 'or' is processed by dividing by 'and'
     for keyword in keyword_divided_or.split(CONFIG.AND_SEARCH_CHARACTER):
-        key = f"{keyword}_{hint['name']}"
+        key = f"{keyword}_{hint.name}"
 
         # Skip if keywords overlap
         if key in duplicate_keys:
@@ -742,29 +740,29 @@ def _build_queries_along_keywords(
     """
 
     # Get the keyword.
-    hints = [x for x in hint_attrs if "keyword" in x and x["keyword"]]
+    hints = [x for x in hint_attrs if x.keyword]
     res_query: dict[str, Any] = {}
 
     for hint in hints:
         and_query: dict[str, Any] = {}
         or_query: dict[str, Any] = {}
 
         # Split keyword by 'or'
-        for keyword_divided_or in hint["keyword"].split(CONFIG.OR_SEARCH_CHARACTER):
+        for keyword_divided_or in (hint.keyword or "").split(CONFIG.OR_SEARCH_CHARACTER):
             if CONFIG.AND_SEARCH_CHARACTER in keyword_divided_or:
                 # If 'AND' is included in the keyword divided by 'OR', add it to 'filter'
                 for keyword in keyword_divided_or.split(CONFIG.AND_SEARCH_CHARACTER):
                     if keyword_divided_or not in and_query:
                         and_query[keyword_divided_or] = {"bool": {"filter": []}}
 
                     and_query[keyword_divided_or]["bool"]["filter"].append(
-                        attr_query[keyword + "_" + hint["name"]]
+                        attr_query[keyword + "_" + hint.name]
                     )
 
             else:
-                and_query[keyword_divided_or] = attr_query[keyword_divided_or + "_" + hint["name"]]
+                and_query[keyword_divided_or] = attr_query[keyword_divided_or + "_" + hint.name]
 
-            if CONFIG.OR_SEARCH_CHARACTER in hint["keyword"]:
+            if CONFIG.OR_SEARCH_CHARACTER in (hint.keyword or ""):
                 # If the keyword contains 'or', concatenate with 'should'
                 if not or_query:
                     or_query = {"bool": {"should": []}}
@@ -817,7 +815,7 @@ def _make_an_attribute_filter(hint: AttrHint, keyword: str) -> dict[str, dict]:
         dict[str, str]: Created attribute filter
 
     """
-    cond_attr: list[dict] = [{"term": {"attr.name": hint["name"]}}]
+    cond_attr: list[dict] = [{"term": {"attr.name": hint.name}}]
 
     date_results = _is_date(keyword)
     if date_results:
@@ -840,7 +838,7 @@ def _make_an_attribute_filter(hint: AttrHint, keyword: str) -> dict[str, dict]:
 
         str_cond = {"regexp": {"attr.value": _get_regex_pattern(keyword)}}
 
-        if hint.get("filter_key") == FilterKey.TEXT_NOT_CONTAINED:
+        if hint.filter_key == FilterKey.TEXT_NOT_CONTAINED:
             cond_attr.append({"bool": {"must_not": [date_cond, str_cond]}})
         else:
             cond_attr.append({"bool": {"should": [date_cond, str_cond]}})
@@ -866,10 +864,10 @@ def _make_an_attribute_filter(hint: AttrHint, keyword: str) -> dict[str, dict]:
             )
 
         elif hint_keyword_val:
-            if "exact_match" not in hint:
+            if hint.exact_match is None:
                 cond_val.append({"regexp": {"attr.value": _get_regex_pattern(hint_keyword_val)}})
 
-            if hint.get("filter_key") == FilterKey.TEXT_NOT_CONTAINED:
+            if hint.filter_key == FilterKey.TEXT_NOT_CONTAINED:
                 cond_attr.append({"bool": {"must_not": cond_val}})
             else:
                 cond_attr.append({"bool": {"should": cond_val}})
@@ -1004,7 +1002,7 @@ def make_search_results(
         # formalize attribute values according to the type
         for attrinfo in entry_info["attr"]:
             # Skip other than the target Attribute
-            if attrinfo["name"] not in [x["name"] for x in hint_attrs]:
+            if attrinfo["name"] not in [x.name for x in hint_attrs]:
                 continue
 
             ret_attrinfo: AdvancedSearchResultRecordAttr = {}
@@ -1024,7 +1022,7 @@ def make_search_results(
                     record.attrs[attrinfo["name"]] = ret_attrinfo
 
             # Check for has permission to EntityAttr
-            if attrinfo["name"] not in [x["name"] for x in hint_attrs if x["is_readable"]]:
+            if attrinfo["name"] not in [x.name for x in hint_attrs if x.is_readable]:
                 ret_attrinfo["is_readable"] = False
                 continue
 

airone/tests/test_elasticsearch.py

@@ -1,7 +1,7 @@
 from django.test import TestCase
 
 from airone.lib import elasticsearch
-from airone.lib.elasticsearch import AdvancedSearchResultRecord
+from airone.lib.elasticsearch import AdvancedSearchResultRecord, AttrHint
 from airone.lib.types import AttrType
 from entity.models import Entity, EntityAttr
 from entry.models import Attribute, AttributeValue, Entry
@@ -42,8 +42,8 @@ def test_make_query(self):
         query = elasticsearch.make_query(
             hint_entity=self._entity,
             hint_attrs=[
-                {"name": "a1", "keyword": "hoge|fu&ga"},
-                {"name": "a2", "keyword": ""},
+                AttrHint(name="a1", keyword="hoge|fu&ga"),
+                AttrHint(name="a2", keyword=""),
             ],
             entry_name="entry1",
         )
@@ -359,7 +359,9 @@ def test_make_search_results(self):
             }
         }
 
-        hint_attrs = [{"name": "test_attr", "keyword": "", "is_readable": True}]
+        hint_attrs = [
+            AttrHint(name="test_attr", keyword="", is_readable=True),
+        ]
         hint_referral = ""
         results = elasticsearch.make_search_results(self._user, res, hint_attrs, hint_referral, 100)
 

api_v1/entry/serializer.py

@@ -2,6 +2,7 @@
 from rest_framework.exceptions import ValidationError
 
 from airone.exceptions import ElasticsearchException
+from airone.lib.elasticsearch import AttrHint
 from airone.lib.log import Logger
 from entity.models import Entity, EntityAttr
 from entry.models import Entry
@@ -294,10 +295,10 @@ def _do_forward_search(sub_query, sub_query_result):
 
             # Query for forward search
             hint_attrs = [
-                {
-                    "name": sub_query["name"],
-                    "keyword": search_keyword,
-                }
+                AttrHint(
+                    name=sub_query["name"],
+                    keyword=search_keyword,
+                )
             ]
 
             # get Entry informations from result

api_v1/entry/views.py

@@ -3,12 +3,14 @@
 import pytz
 from django.conf import settings
 from django.db.models import Q
+from pydantic import ValidationError
 from rest_framework import status
 from rest_framework.response import Response
 from rest_framework.views import APIView
 
 from airone.exceptions import ElasticsearchException
 from airone.lib.acl import ACLType
+from airone.lib.elasticsearch import AttrHint
 from api_v1.entry.serializer import EntrySearchChainSerializer
 from entity.models import Entity
 from entry.models import Entry
@@ -62,15 +64,15 @@ def post(self, request, format=None):
 
         hint_entities = request.data.get("entities")
         hint_entry_name = request.data.get("entry_name", "")
-        hint_attrs = request.data.get("attrinfo")
         hint_referral = request.data.get("referral")
+        attrinfo = request.data.get("attrinfo", [])
         is_output_all = request.data.get("is_output_all", True)
         entry_limit = request.data.get("entry_limit", CONFIG_ENTRY.MAX_LIST_ENTRIES)
 
         if (
             not isinstance(hint_entities, list)
             or not isinstance(hint_entry_name, str)
-            or not isinstance(hint_attrs, list)
+            or not isinstance(attrinfo, list)
             or not isinstance(is_output_all, bool)
             or (hint_referral and not isinstance(hint_referral, str))
             or not isinstance(entry_limit, int)
@@ -79,21 +81,31 @@ def post(self, request, format=None):
                 "The type of parameter is incorrect", status=status.HTTP_400_BAD_REQUEST
             )
 
+        try:
+            hint_attrs = [
+                AttrHint(
+                    name=x.get("name"),
+                    keyword=x.get("keyword"),
+                    filter_key=x.get("filter_key"),
+                    exact_match=x.get("exact_match"),
+                )
+                for x in attrinfo
+            ]
+        except (TypeError, ValidationError):
+            return Response(
+                "The type of parameter 'attrinfo' is incorrect",
+                status=status.HTTP_400_BAD_REQUEST,
+            )
+
         # forbid to input large size request
         if len(hint_entry_name) > CONFIG_ENTRY.MAX_QUERY_SIZE:
             return Response("Sending parameter is too large", status=400)
 
         # check attribute params
         for hint_attr in hint_attrs:
-            if "name" not in hint_attr:
-                return Response("The name key is required for attrinfo parameter", status=400)
-            if not isinstance(hint_attr["name"], str):
-                return Response("Invalid value for attrinfo parameter", status=400)
-            if hint_attr.get("keyword"):
-                if not isinstance(hint_attr["keyword"], str):
-                    return Response("Invalid value for attrinfo parameter", status=400)
+            if hint_attr.keyword:
                 # forbid to input large size request
-                if len(hint_attr["keyword"]) > CONFIG_ENTRY.MAX_QUERY_SIZE:
+                if len(hint_attr.keyword) > CONFIG_ENTRY.MAX_QUERY_SIZE:
                     return Response("Sending parameter is too large", status=400)
 
         # check entities params

api_v1/tests/entry/test_api.py

@@ -29,25 +29,24 @@ def test_search_invalid_param(self):
             params = {**valid_params, **invalid_param}
             resp = self.client.post("/api/v1/entry/search", json.dumps(params), "application/json")
             self.assertEqual(resp.status_code, 400)
-            self.assertEqual(resp.content, b'"The type of parameter is incorrect"')
 
         params = {**valid_params, **{"attrinfo": [{"hoge": "value"}]}}
         resp = self.client.post("/api/v1/entry/search", json.dumps(params), "application/json")
         self.assertEqual(resp.status_code, 400)
-        self.assertEqual(resp.content, b'"The name key is required for attrinfo parameter"')
+        self.assertEqual(resp.content, b"\"The type of parameter 'attrinfo' is incorrect\"")
 
         params = {**valid_params, **{"attrinfo": [{"name": ["hoge"]}]}}
         resp = self.client.post("/api/v1/entry/search", json.dumps(params), "application/json")
         self.assertEqual(resp.status_code, 400)
-        self.assertEqual(resp.content, b'"Invalid value for attrinfo parameter"')
+        self.assertEqual(resp.content, b"\"The type of parameter 'attrinfo' is incorrect\"")
 
         params = {
             **valid_params,
             **{"attrinfo": [{"name": "value", "keyword": ["hoge"]}]},
         }
         resp = self.client.post("/api/v1/entry/search", json.dumps(params), "application/json")
         self.assertEqual(resp.status_code, 400)
-        self.assertEqual(resp.content, b'"Invalid value for attrinfo parameter"')
+        self.assertEqual(resp.content, b"\"The type of parameter 'attrinfo' is incorrect\"")
 
     def test_narrow_down_advanced_search_results(self):
         user = self.admin_login()

dashboard/tasks.py

@@ -8,7 +8,7 @@
 from natsort import natsorted
 
 from airone.celery import app
-from airone.lib.elasticsearch import AdvancedSearchResultRecord
+from airone.lib.elasticsearch import AdvancedSearchResultRecord, AttrHint
 from airone.lib.job import may_schedule_until_job_is_ready
 from airone.lib.types import AttrType
 from entry.services import AdvancedSearchService
@@ -174,13 +174,14 @@ def export_search_result(self, job: Job):
     has_referral: bool = recv_data.get("has_referral", False)
     referral_name: str | None = recv_data.get("referral_name")
     entry_name: str | None = recv_data.get("entry_name")
+    hint_attrs = [AttrHint.model_validate(attr) for attr in recv_data["attrinfo"]]
     if has_referral and referral_name is None:
         referral_name = ""
 
     resp = AdvancedSearchService.search_entries(
         user,
         recv_data["entities"],
-        recv_data["attrinfo"],
+        hint_attrs,
         settings.ES_CONFIG["MAXIMUM_RESULTS_NUM"],
         entry_name,
         referral_name,

dashboard/tests/test_view.py

@@ -213,7 +213,16 @@ def test_show_advanced_search_results(self):
         attr = entry.attrs.first()
         self.assertEqual(resp.status_code, 200)
         self.assertEqual(
-            resp.context["hint_attrs"], [{"name": "attr", "is_readable": attr.is_public}]
+            resp.context["hint_attrs"],
+            [
+                {
+                    "name": "attr",
+                    "is_readable": attr.is_public,
+                    "filter_key": None,
+                    "keyword": None,
+                    "exact_match": None,
+                }
+            ],
         )
         self.assertEqual(resp.context["results"]["ret_count"], 20)
         self.assertEqual(len(resp.context["results"]["ret_values"]), 20)
@@ -320,7 +329,7 @@ def test_show_advanced_search_results(self):
         self.assertEqual(resp.status_code, 400)
         self.assertEqual(
             resp.content.decode("utf-8"),
-            "The name key is required for attrinfo parameter",
+            "Invalid value for attrinfo parameter",
         )
 
         resp = self.client.get(