content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import json
  import logging
  from typing import (
      Any, Dict, List,
  )
  
  from amundsen_common.models.search import (
      Filter, HighlightOptions, SearchResponse,
  )
  from elasticsearch import Elasticsearch
  from elasticsearch_dsl import (
      MultiSearch, Q, Search,
  )
  from elasticsearch_dsl.query import Match, RankFeature
  from elasticsearch_dsl.response import Response
  from flask import current_app
  
  from search_service import config
  from search_service.proxy.es_proxy_utils import Resource, create_search_response
  from search_service.proxy.es_proxy_v2 import BOOL_QUERY, ElasticsearchProxyV2
  
  LOGGER = logging.getLogger(__name__)
  
  # ES query constants
  
  DEFAULT_FUZZINESS = "AUTO"
  
  # using fvh as the default highlighter because it supports very large documents
  DEFAULT_HIGHLIGHTER = 'fvh'
  
  
  class ElasticsearchProxyV2_1(ElasticsearchProxyV2):
  
      # map the field name in FE to the field used to filter in ES
      # note: ES needs keyword field types to filter
  
      GENERAL_MAPPING = {
          'key': 'key',
          'description': 'description',
          'resource_type': 'resource_type',
      }
  
      TABLE_MAPPING = {
          **GENERAL_MAPPING,
          'badges': 'badges.keyword',
          'tag': 'tags.keyword',
          'schema': 'schema.keyword',
          'table': 'name.keyword',
          'column': 'columns.keyword',
          'database': 'database.keyword',
          'cluster': 'cluster.keyword',
      }
  
      DASHBOARD_MAPPING = {
          **GENERAL_MAPPING,
          'url': 'url',
          'uri': 'uri',
          'last_successful_run_timestamp': 'last_successful_run_timestamp',
          'group_name': 'group_name.keyword',
          'chart_names': 'chart_names.keyword',
          'query_names': 'query_names.keyword',
          'name': 'name.keyword',
          'tag': 'tags.keyword',
      }
  
      FEATURE_MAPPING = {
          **GENERAL_MAPPING,
          'version': 'version',
          'availability': 'availability',
          'feature_group': 'feature_group.keyword',
          'feature_name': 'name.keyword',
          'entity': 'entity.keyword',
          'status': 'status.keyword',
          'tags': 'tags.keyword',
          'badges': 'badges.keyword',
      }
  
      USER_MAPPING = {
          'full_name': 'name.keyword',
          'email': 'key',
          'first_name': 'first_name',
          'last_name': 'last_name',
          'resource_type': 'resource_type',
      }
  
      RESOURCE_TO_MAPPING = {
          Resource.TABLE: TABLE_MAPPING,
          Resource.DASHBOARD: DASHBOARD_MAPPING,
          Resource.FEATURE: FEATURE_MAPPING,
          Resource.USER: USER_MAPPING,
      }
  
      # The overriding of __new__ here is a temporary solution to provide backwards compatiblity
      # until most of the community has moved to using the new Elasticsearch mappings and it will
      # be removed once ElasticsearchProxyV2 id deprecated
      def __new__(cls: Any,
                  host: str,
                  user: str,
                  password: str,
                  client: Elasticsearch,
                  page_size: int, *args: str, **kwargs: int) -> Any:
          elasticsearch_client = None
          if client:
              elasticsearch_client = client
          else:
              http_auth = (user, password) if user else None
              elasticsearch_client = Elasticsearch(host, http_auth=http_auth)
  
          # check if any index uses the most up to date mappings (version == 2)
          indices = elasticsearch_client.indices.get_alias(index='*')
          mappings_up_to_date = False
          for index in indices:
              index_mapping = elasticsearch_client.indices.get_mapping(index=index).get(index)
              mapping_meta_field = index_mapping.get('mappings').get('_meta')
              if mapping_meta_field is not None and mapping_meta_field.get('version') == 2:
                  mappings_up_to_date = True
                  break
  
          if mappings_up_to_date:
              # Use ElasticsearchProxyV2_1 if indexes are up to date with mappings
              obj = super(ElasticsearchProxyV2_1, cls).__new__(cls)
              return obj
  
          # If old mappings are used proxy client should be ElasticsearchProxyV2
          obj = super(ElasticsearchProxyV2_1, cls).__new__(ElasticsearchProxyV2)
          obj.__init__(host=host,
                       user=user,
                       password=password,
                       client=elasticsearch_client,
                       page_size=page_size)
          return obj
  
      def get_index_alias_for_resource(self, resource_type: Resource) -> str:
          resource_str = resource_type.name.lower()
          alias_config = current_app.config.get(
              config.ES_INDEX_ALIAS_TEMPLATE
          )
          if alias_config is None:
              return f'{resource_str}_search_index_v2_1'
  
          alias = str(alias_config).format(resource=resource_str)
  
          return alias
  
      def _build_must_query(self, resource: Resource, query_term: str) -> List[Q]:
          """
          Builds the query object for the inputed search term
          """
  
          if not query_term:
              # We don't want to create match query for ""
              # because it will result in no matches even with filters
              return []
  
          # query for fields general to all resources
          should_clauses: List[Q] = [
              Match(name={
                  "query": query_term,
                  "fuzziness": DEFAULT_FUZZINESS,
                  "boost": 5
              }),
              Match(description={
                  "query": query_term,
                  "fuzziness": DEFAULT_FUZZINESS,
                  "boost": 1.5
              }),
              Match(badges={
                  "query": query_term,
                  "fuzziness": DEFAULT_FUZZINESS,
              }),
              Match(tags={
                  "query": query_term,
                  "fuzziness": DEFAULT_FUZZINESS,
              }),
          ]
  
          if resource == Resource.TABLE:
              columns_subfield = 'columns.general'
              should_clauses.extend([
                  Match(schema={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 3
                  }),
                  Match(**{columns_subfield: {
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 2
                  }}),
                  Match(column_descriptions={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS
                  }),
              ])
          elif resource == Resource.DASHBOARD:
              should_clauses.extend([
                  Match(group_name={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 3
                  }),
                  Match(query_names={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 2
                  }),
                  Match(chart_names={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 2
                  }),
                  Match(uri={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 4
                  }),
              ])
          elif resource == Resource.FEATURE:
              should_clauses.extend([
                  Match(feature_group={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 3
                  }),
                  Match(version={
                      "query": query_term
                  }),
                  Match(entity={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 2
                  }),
                  Match(status={
                      "query": query_term
                  }),
              ])
          elif resource == Resource.USER:
              # replaces rather than extending
              should_clauses = [
                  Match(name={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 5
                  }),
                  Match(first_name={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 3
                  }),
                  Match(last_name={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 3
                  }),
                  Match(team_name={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS
                  }),
                  Match(key={
                      "query": query_term,
                      "fuzziness": DEFAULT_FUZZINESS,
                      "boost": 4
                  }),
              ]
  
          must_clauses: List[Q] = [Q(BOOL_QUERY, should=should_clauses)]
  
          return must_clauses
  
      def _build_should_query(self, resource: Resource, query_term: str) -> List[Q]:
  
          # no scoring happens if there is no search term
          if query_term == '':
              return []
  
          # general usage metric for searcheable resources
          usage_metric_field_boosts = {
              'total_usage': 10.0,
          }
  
          if resource == Resource.TABLE:
              usage_metric_field_boosts = {
                  **usage_metric_field_boosts,
                  'unique_usage': 10.0,
              }
          if resource == Resource.USER:
              usage_metric_field_boosts = {
                  'total_read': 10.0,
                  'total_own': 10.0,
                  'total_follow': 10.0,
              }
  
          rank_feature_queries = []
  
          for metric in usage_metric_field_boosts.keys():
              field_name = f'usage.{metric}'
              boost = usage_metric_field_boosts[metric]
              rank_feature_query = RankFeature(field=field_name,
                                               boost=boost)
              rank_feature_queries.append(rank_feature_query)
  
          return rank_feature_queries
  
      def _search_highlight(self,
                            resource: Resource,
                            search: Search,
                            highlight_options: Dict[Resource, HighlightOptions]) -> Search:
          # get highlighting options for resource
          resource_options = highlight_options.get(resource)
          highlighting_enabled = resource_options.enable_highlight if resource_options else None
  
          if highlighting_enabled:
              # default highlighted fields
              search = search.highlight('name',
                                        type=DEFAULT_HIGHLIGHTER,
                                        number_of_fragments=0)
              search = search.highlight('description',
                                        type=DEFAULT_HIGHLIGHTER,
                                        number_of_fragments=0)
              if resource == Resource.TABLE:
                  search = search.highlight('columns.general',
                                            type=DEFAULT_HIGHLIGHTER,
                                            number_of_fragments=10,
                                            order='score')
                  search = search.highlight('column_descriptions',
                                            type=DEFAULT_HIGHLIGHTER,
                                            number_of_fragments=5,
                                            order='score')
              if resource == Resource.DASHBOARD:
                  search = search.highlight('chart_names',
                                            type=DEFAULT_HIGHLIGHTER,
                                            number_of_fragments=10,
                                            order='score')
                  search = search.highlight('query_names',
                                            type=DEFAULT_HIGHLIGHTER,
                                            number_of_fragments=10,
                                            order='score')
  
          return search
  
      def execute_multisearch_query(self, multisearch: MultiSearch) -> List[Response]:
          try:
              response = multisearch.execute()
              return response
          except Exception as e:
              LOGGER.error(f'Failed to execute ES search queries. {e}')
              return []
  
      def search(self, *,
                 query_term: str,
                 page_index: int,
                 results_per_page: int,
                 resource_types: List[Resource],
                 filters: List[Filter],
                 highlight_options: Dict[Resource, HighlightOptions]) -> SearchResponse:
          if not resource_types:
              # if resource types are not defined then search all resources
              resource_types = self.PRIMARY_ENTITIES
  
          multisearch = MultiSearch(using=self.elasticsearch)
  
          for resource in resource_types:
              # guard clause to prevent search in missing indices or aliases
              aliases_in_es = {alias['alias'] for alias in self.elasticsearch.cat.aliases(format="json")}
              resource_alias = self.get_index_alias_for_resource(resource_type=resource)
              if resource_alias not in aliases_in_es:
                  LOGGER.info(f"There are no indices in elasticsearch against resource_type: {resource}")
                  continue
  
              # build a query for each resource to search
              query_for_resource = self._build_elasticsearch_query(resource=resource,
                                                                   query_term=query_term,
                                                                   filters=filters)
              # wrap the query in a search object
              search = Search(index=resource_alias).query(query_for_resource)
  
              # highlighting
              if highlight_options:
                  search = self._search_highlight(resource=resource,
                                                  search=search,
                                                  highlight_options=highlight_options)
  
              # pagination
              start_from = page_index * results_per_page
              end = results_per_page * (page_index + 1)
              search = search[start_from:end]
  
              # add search object to multisearch
              LOGGER.info(json.dumps(search.to_dict()))
              multisearch = multisearch.add(search)
  
          responses = self.execute_multisearch_query(multisearch=multisearch)
  
          formatted_response = create_search_response(page_index=page_index,
                                                      results_per_page=results_per_page,
                                                      responses=responses,
                                                      resource_types=resource_types,
                                                      resource_to_field_mapping=self.RESOURCE_TO_MAPPING)
  
          return formatted_response
  