content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import json
  import logging
  from typing import (
      Any, Dict, List, Optional, Union,
  )
  
  from amundsen_common.models.api import health_check
  from amundsen_common.models.search import (
      Filter, HighlightOptions, SearchResponse,
  )
  from elasticsearch import Elasticsearch
  from elasticsearch.exceptions import ConnectionError as ElasticConnectionError, ElasticsearchException
  from elasticsearch_dsl import (
      MultiSearch, Q, Search,
  )
  from elasticsearch_dsl.query import MultiMatch
  from elasticsearch_dsl.response import Response
  from elasticsearch_dsl.utils import AttrList
  from werkzeug.exceptions import InternalServerError
  
  from search_service.proxy.es_proxy_utils import Resource, create_search_response
  
  LOGGER = logging.getLogger(__name__)
  
  BOOL_QUERY = 'bool'
  WILDCARD_QUERY = 'wildcard'
  TERM_QUERY = 'term'
  TERMS_QUERY = 'terms'
  
  
  class ElasticsearchProxyV2():
      PRIMARY_ENTITIES = [Resource.TABLE, Resource.DASHBOARD, Resource.FEATURE, Resource.USER]
  
      # mapping to translate request for table resources
      TABLE_MAPPING = {
          'key': 'key',
          'badges': 'badges',
          'tag': 'tags',
          'schema': 'schema.raw',
          'table': 'name.raw',
          'column': 'column_names.raw',
          'database': 'database.raw',
          'cluster': 'cluster.raw',
          'description': 'description',
          'resource_type': 'resource_type'
      }
  
      # mapping to translate request for dashboard resources
      DASHBOARD_MAPPING = {
          'group_name': 'group_name.raw',
          'group_url': 'group_url',
          'url': 'url',
          'uri': 'uri',
          'name': 'name.raw',
          'product': 'product',
          'tag': 'tags',
          'description': 'description',
          'last_successful_run_timestamp': 'last_successful_run_timestamp',
          'resource_type': 'resource_type'
      }
  
      # mapping to translate request for feature resources
      FEATURE_MAPPING = {
          'key': 'key',
          'feature_group': 'feature_group.raw',
          'feature_name': 'feature_name.raw',
          'entity': 'entity',
          'status': 'status',
          'version': 'version',
          'availability': 'availability.raw',
          'tags': 'tags',
          'badges': 'badges',
          'description': 'description',
          'resource_type': 'resource_type'
      }
  
      USER_MAPPING = {
          'full_name': 'full_name',
          'first_name': 'first_name',
          'last_name': 'last_name',
          'email': 'email',
          'resource_type': 'resource_type'
      }
  
      RESOURCE_TO_MAPPING = {
          Resource.TABLE: TABLE_MAPPING,
          Resource.DASHBOARD: DASHBOARD_MAPPING,
          Resource.FEATURE: FEATURE_MAPPING,
          Resource.USER: USER_MAPPING,
      }
  
      MUST_FIELDS_TABLE = ['name^3',
                           'name.raw^3',
                           'schema^2',
                           'description',
                           'column_names',
                           'badges']
  
      MUST_FIELDS_DASHBOARD = ['name.raw^75',
                               'name^7',
                               'group_name.raw^15',
                               'group_name^7',
                               'description^3',
                               'query_names^3',
                               'chart_names^2']
  
      MUST_FIELDS_FEATURE = ['feature_name.raw^25',
                             'feature_name^7',
                             'feature_group.raw^15',
                             'feature_group^7',
                             'version^7',
                             'description^3',
                             'status',
                             'entity',
                             'tags',
                             'badges']
  
      MUST_FIELDS_USER = ['full_name.raw^30',
                          'full_name^5',
                          'first_name.raw^5',
                          'last_name.raw^5',
                          'first_name^3',
                          'last_name^3',
                          'email^3']
  
      def __init__(self, *,
                   host: Optional[str] = None,
                   user: str = '',
                   password: str = '',
                   client: Optional[Elasticsearch] = None,
                   page_size: int = 10) -> None:
          if client:
              self.elasticsearch = client
          else:
              http_auth = (user, password) if user else None
              self.elasticsearch = Elasticsearch(host, http_auth=http_auth)
          self.page_size = page_size
  
      def health(self) -> health_check.HealthCheck:
          """
          Returns the health of the Elastic search cluster
          """
          try:
              if self.elasticsearch.ping():
                  health = self.elasticsearch.cluster.health()
                  # ES status vaues: green, yellow, red
                  status = health_check.OK if health['status'] != 'red' else health_check.FAIL
              else:
                  health = {'status': 'Unable to connect'}
                  status = health_check.FAIL
              checks = {f'{type(self).__name__}:connection': health}
          except ElasticConnectionError:
              status = health_check.FAIL
              checks = {f'{type(self).__name__}:connection': {'status': 'Unable to connect'}}
          return health_check.HealthCheck(status=status, checks=checks)
  
      def _get_must_fields(self, resource: Resource) -> List[str]:
          must_fields_mapping = {
              Resource.TABLE: self.MUST_FIELDS_TABLE,
              Resource.DASHBOARD: self.MUST_FIELDS_DASHBOARD,
              Resource.FEATURE: self.MUST_FIELDS_FEATURE,
              Resource.USER: self.MUST_FIELDS_USER
          }
  
          return must_fields_mapping[resource]
  
      def get_index_alias_for_resource(self, resource_type: Resource) -> str:
          resource_str = resource_type.name.lower()
          return f'{resource_str}_search_index'
  
      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 multi_match query for ""
              # because it will result in no matches even with filters
              return []
  
          try:
              fields: List[str] = self._get_must_fields(resource)
          except KeyError:
              # TODO if you don't specify a resource match for all generic fields in the future
              raise ValueError(f'no fields defined for resource {resource}')
  
          return [MultiMatch(query=query_term, fields=fields, type='cross_fields')]
  
      def _build_should_query(self, resource: Resource, query_term: str) -> List[Q]:
          # Can define on custom es_search_proxy class, no default implementation
          return []
  
      def _build_filters(self, resource: Resource, filters: List[Filter]) -> List:
          """
          Builds the query object for all of the filters given in the search request
          """
          mapping = type(self).RESOURCE_TO_MAPPING.get(resource)
  
          filter_queries: List = []
  
          for filter in filters:
              if mapping is not None and mapping.get(filter.name) is not None:
                  # only apply filter to query if field exists for the given resource
                  filter_name = mapping.get(filter.name)
  
                  queries_per_term = [Q(WILDCARD_QUERY, **{filter_name: term}) for term in filter.values]
  
                  if filter.operation == 'OR':
                      filter_queries.append(Q(BOOL_QUERY, should=queries_per_term, minimum_should_match=1))
  
                  elif filter.operation == 'AND':
                      for q in queries_per_term:
                          filter_queries.append(q)
  
                  else:
                      msg = f'Invalid operation {filter.operation} for filter {filter_name} with values {filter.values}'
                      raise ValueError(msg)
              else:
                  LOGGER.info(f'Filter {filter.name} does not apply to {resource}')
  
          return filter_queries
  
      def _build_elasticsearch_query(self, *,
                                     resource: Resource,
                                     query_term: str,
                                     filters: List[Filter]) -> Q:
  
          must_query = self._build_must_query(resource=resource,
                                              query_term=query_term)
  
          should_query = self._build_should_query(resource=resource,
                                                  query_term=query_term)
  
          filters = self._build_filters(resource=resource, filters=filters)
  
          es_query = Q(BOOL_QUERY, must=must_query, should=should_query, filter=filters)
  
          return es_query
  
      def execute_queries(self, queries: Dict[Resource, Q],
                          page_index: int,
                          results_per_page: int) -> List[Response]:
          multisearch = MultiSearch(using=self.elasticsearch)
  
          for resource in queries.keys():
              query_for_resource = queries.get(resource)
              search = Search(index=self.get_index_alias_for_resource(resource_type=resource)).query(query_for_resource)
              LOGGER.info(json.dumps(search.to_dict()))
  
              # pagination
              start_from = page_index * results_per_page
              end = results_per_page * (page_index + 1)
  
              search = search[start_from:end]
  
              multisearch = multisearch.add(search)
          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
  
          queries: Dict[Resource, Q] = {}
          for resource in resource_types:
              # build a query for each resource to search
              queries[resource] = self._build_elasticsearch_query(resource=resource,
                                                                  query_term=query_term,
                                                                  filters=filters)
  
          responses = self.execute_queries(queries=queries,
                                           page_index=page_index,
                                           results_per_page=results_per_page)
  
          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
  
      def get_document_json_by_key(self,
                                   resource_key: str,
                                   resource_type: Resource) -> Any:
          key_query = {
              resource_type: Q(TERM_QUERY, key=resource_key),
          }
          response: Response = self.execute_queries(queries=key_query,
                                                    page_index=0,
                                                    results_per_page=1)
          if len(response) < 1:
              msg = f'No response from ES for key query {key_query[resource_type]}'
              LOGGER.error(msg)
              raise ElasticsearchException(msg)
  
          response = response[0]
          if response.success():
              results_count = response.hits.total.value
              if results_count == 1:
                  es_result = response.hits.hits[0]
                  return es_result
  
              if results_count > 1:
                  msg = f'Key {key_query[resource_type]} is not unique to a single ES resource'
                  LOGGER.error(msg)
                  raise ValueError(msg)
  
              else:
                  # no doc exists with given key in ES
                  msg = f'Requested key {resource_key} query returned no results in ES'
                  LOGGER.error(msg)
                  raise ValueError(msg)
          else:
              msg = f'Request to Elasticsearch failed: {response.failures}'
              LOGGER.error(msg)
              raise InternalServerError(msg)
  
      def update_document_by_id(self, *,
                                resource_type: Resource,
                                field: str,
                                new_value: Union[List, str, None],
                                document_id: str) -> None:
  
          partial_document = {
              'doc': {
                  field: new_value
              }
          }
          self.elasticsearch.update(index=self.get_index_alias_for_resource(resource_type=resource_type),
                                    id=document_id,
                                    body=partial_document)
  
      def update_document_by_key(self, *,
                                 resource_key: str,
                                 resource_type: Resource,
                                 field: str,
                                 value: Optional[str] = None,
                                 operation: str = 'add') -> str:
  
          mapped_field = self.RESOURCE_TO_MAPPING[resource_type].get(field)
          if not mapped_field:
              mapped_field = field
  
          try:
              es_hit = self.get_document_json_by_key(resource_key=resource_key,
                                                     resource_type=resource_type)
              document_id = es_hit._id
              current_value = getattr(es_hit._source, mapped_field)
  
          except Exception as e:
              msg = f'Failed to get ES document id and current value for key {resource_key}. {e}'
              LOGGER.error(msg)
              return msg
  
          new_value = current_value
  
          if operation == 'overwrite':
              if type(current_value) is AttrList:
                  new_value = [value]
              else:
                  new_value = value
          else:
              # operation is add
              if type(current_value) is AttrList:
                  curr_list = list(current_value)
                  curr_list.append(value)
                  new_value = curr_list
              else:
                  new_value = [current_value, value]
  
          try:
              self.update_document_by_id(resource_type=resource_type,
                                         field=mapped_field,
                                         new_value=new_value,
                                         document_id=document_id)
          except Exception as e:
              msg = f'Failed to update field {field} with value {new_value} for {resource_key}. {e}'
              LOGGER.error(msg)
              return msg
  
          return f'ES document field {field} for {resource_key} with value {value} was updated successfully'
  
      def delete_document_by_key(self, *,
                                 resource_key: str,
                                 resource_type: Resource,
                                 field: str,
                                 value: Optional[str] = None) -> str:
          mapped_field = self.RESOURCE_TO_MAPPING[resource_type].get(field)
          if not mapped_field:
              mapped_field = field
  
          try:
              es_hit = self.get_document_json_by_key(resource_key=resource_key,
                                                     resource_type=resource_type)
              document_id = es_hit._id
              current_value = getattr(es_hit._source, mapped_field)
  
          except Exception as e:
              msg = f'Failed to get ES document id and current value for key {resource_key}. {e}'
              LOGGER.error(msg)
              return msg
  
          new_value = current_value
  
          if type(current_value) is AttrList:
              if value:
                  curr_list = list(current_value)
                  curr_list.remove(value)
                  new_value = curr_list
              else:
                  new_value = []
          else:
              # no value given when deleting implies
              # delete is happening on a single value field
              new_value = ''
          try:
              self.update_document_by_id(resource_type=resource_type,
                                         field=mapped_field,
                                         new_value=new_value,
                                         document_id=document_id)
          except Exception as e:
              msg = f'Failed to delete field {field} with value {new_value} for {resource_key}. {e}'
              LOGGER.error(msg)
              return msg
  
          return f'ES document field {field} for {resource_key} with value {value} was deleted successfully'
  