content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import importlib
  from typing import (
      Any, Dict, List, Optional,
  )
  
  from gremlin_python.process.graph_traversal import GraphTraversalSource, __
  from gremlin_python.process.traversal import (
      Order, T, TextP,
  )
  from pyhocon import ConfigTree
  
  from databuilder import Scoped
  from databuilder.clients.neptune_client import NeptuneSessionClient
  from databuilder.extractor.base_extractor import Extractor
  from databuilder.models.cluster.cluster_constants import CLUSTER_REVERSE_RELATION_TYPE
  from databuilder.models.dashboard.dashboard_chart import DashboardChart
  from databuilder.models.dashboard.dashboard_metadata import DashboardMetadata
  from databuilder.models.dashboard.dashboard_query import DashboardQuery
  from databuilder.models.description_metadata import DescriptionMetadata
  from databuilder.models.owner_constants import OWNER_OF_OBJECT_RELATION_TYPE
  from databuilder.models.schema.schema_constant import SCHEMA_REVERSE_RELATION_TYPE
  from databuilder.models.table_metadata import TableMetadata
  from databuilder.models.timestamp.timestamp_constants import LASTUPDATED_RELATION_TYPE, TIMESTAMP_PROPERTY
  from databuilder.models.usage.usage_constants import (
      READ_RELATION_COUNT_PROPERTY, READ_RELATION_TYPE, READ_REVERSE_RELATION_TYPE,
  )
  from databuilder.models.user import User
  from databuilder.serializers.neptune_serializer import METADATA_KEY_PROPERTY_NAME
  
  
  def _table_search_query(graph: GraphTraversalSource, tag_filter: str) -> List[Dict]:
      traversal = graph.V().hasLabel(TableMetadata.TABLE_NODE_LABEL)
      if tag_filter:
          traversal = traversal.has('published_tag', tag_filter)
      traversal = traversal.project(
          'database',
          'cluster',
          'schema',
          'schema_description',
          'name',
          'key',
          'description',
          'last_updated_timestamp',
          'column_names',
          'column_descriptions',
          'total_usage',
          'unique_usage',
          'tags',
          'badges',
          'programmatic_descriptions'
      )
      traversal = traversal.by(
          __.out(
              TableMetadata.TABLE_SCHEMA_RELATION_TYPE
          ).out(SCHEMA_REVERSE_RELATION_TYPE).out(CLUSTER_REVERSE_RELATION_TYPE).values('name')
      )  # database
      traversal = traversal.by(
          __.out(TableMetadata.TABLE_SCHEMA_RELATION_TYPE).out(SCHEMA_REVERSE_RELATION_TYPE).values('name')
      )  # cluster
      traversal = traversal.by(__.out(TableMetadata.TABLE_SCHEMA_RELATION_TYPE).values('name'))  # schema
      traversal = traversal.by(__.coalesce(
          __.out(TableMetadata.TABLE_SCHEMA_RELATION_TYPE).out(
              DescriptionMetadata.DESCRIPTION_RELATION_TYPE
          ).values('description'),
          __.constant('')
      ))  # schema_description
      traversal = traversal.by('name')  # name
      traversal = traversal.by(T.id)  # key
      traversal = traversal.by(__.coalesce(
          __.out(DescriptionMetadata.DESCRIPTION_RELATION_TYPE).values('description'),
          __.constant('')
      ))  # description
      traversal = traversal.by(
          __.coalesce(__.out(LASTUPDATED_RELATION_TYPE).values(TIMESTAMP_PROPERTY), __.constant(''))
      )  # last_updated_timestamp
      traversal = traversal.by(__.out(TableMetadata.TABLE_COL_RELATION_TYPE).values('name').fold())  # column_names
      traversal = traversal.by(
          __.out(TableMetadata.TABLE_COL_RELATION_TYPE).out(
              DescriptionMetadata.DESCRIPTION_RELATION_TYPE
          ).values('description').fold()
      )  # column_descriptions
      traversal = traversal.by(__.coalesce(
          __.outE(READ_REVERSE_RELATION_TYPE).values('read_count'),
          __.constant(0)).sum()
      )  # total_usage
      traversal = traversal.by(__.outE(READ_REVERSE_RELATION_TYPE).count())  # unique_usage
      traversal = traversal.by(
          __.inE(TableMetadata.TAG_TABLE_RELATION_TYPE).outV().values(METADATA_KEY_PROPERTY_NAME).fold()
      )  # tags
      traversal = traversal.by(
          __.out('HAS_BADGE').values('keys').dedup().fold()
      )  # badges
      traversal = traversal.by(
          __.out(DescriptionMetadata.PROGRAMMATIC_DESCRIPTION_NODE_LABEL).values('description').fold()
      )  # programmatic_descriptions
      traversal = traversal.order().by(__.select('name'), Order.asc)
      return traversal.toList()
  
  
  def _user_search_query(graph: GraphTraversalSource, tag_filter: str) -> List[Dict]:
      traversal = graph.V().hasLabel(User.USER_NODE_LABEL)
      traversal = traversal.has(User.USER_NODE_FULL_NAME)
      if tag_filter:
          traversal = traversal.where('published_tag', tag_filter)
      traversal = traversal.project(
          'email',
          'first_name',
          'last_name',
          'full_name',
          'github_username',
          'team_name',
          'employee_type',
          'manager_email',
          'slack_id',
          'is_active',
          'role_name',
          'total_read',
          'total_own',
          'total_follow'
      )
      traversal = traversal.by('full_name')  # name
      traversal = traversal.by('email')  # email
      traversal = traversal.by('first_name')  # first_name
      traversal = traversal.by('last_name')  # last_name
      traversal = traversal.by('full_name')  # full_name
      traversal = traversal.by('github_username')  # github_username
      traversal = traversal.by('team_name')  # team_name
      traversal = traversal.by('employee_type')  # employee_type
      traversal = traversal.by(__.coalesce(
          __.out(User.USER_MANAGER_RELATION_TYPE).values('email'),
          __.constant(''))
      )  # manager_email
      traversal = traversal.by('slack_id')  # slack_id
      traversal = traversal.by('is_active')  # is_active
      traversal = traversal.by('role_name')  # role_name
      traversal = traversal.by(__.coalesce(
          __.outE(READ_RELATION_TYPE).values('read_count'),
          __.constant(0)
      ).sum())  # total_read
      traversal = traversal.by(__.outE(OWNER_OF_OBJECT_RELATION_TYPE).fold().count())  # total_own
      traversal = traversal.by(__.outE('FOLLOWED_BY').fold().count())  # total_follow
      traversal = traversal.order().by(__.select('email'), Order.asc)
      return traversal.toList()
  
  
  def _dashboard_search_query(graph: GraphTraversalSource, tag_filter: str) -> List[Dict]:
      traversal = graph.V().hasLabel(DashboardMetadata.DASHBOARD_NODE_LABEL)
      traversal = traversal.has('name')
      if tag_filter:
          traversal = traversal.where('published_tag', tag_filter)
  
      traversal = traversal.project(
          'group_name',
          'name',
          'cluster',
          'description',
          'group_description',
          'group_url',
          'url',
          'uri',
          'last_successful_run_timestamp',
          'query_names',
          'chart_names',
          'total_usage',
          'tags',
          'badges'
      )
      traversal = traversal.by(
          __.out(DashboardMetadata.DASHBOARD_DASHBOARD_GROUP_RELATION_TYPE).values('name')
      )  # group_name
      traversal = traversal.by('name')  # name
      traversal = traversal.by(
          __.out(
              DashboardMetadata.DASHBOARD_DASHBOARD_GROUP_RELATION_TYPE
          ).out(
              DashboardMetadata.DASHBOARD_GROUP_CLUSTER_RELATION_TYPE
          ).values('name')
      )  # cluster
      traversal = traversal.by(__.coalesce(
          __.out(DashboardMetadata.DASHBOARD_DESCRIPTION_RELATION_TYPE).values('description'),
          __.constant('')
      ))  # description
      traversal = traversal.by(__.coalesce(
          __.out(DashboardMetadata.DASHBOARD_DASHBOARD_GROUP_RELATION_TYPE).out(
              DashboardMetadata.DASHBOARD_DESCRIPTION_RELATION_TYPE
          ).values('description'),
          __.constant('')
      ))  # group_description
      traversal = traversal.by(
          __.out(DashboardMetadata.DASHBOARD_DASHBOARD_GROUP_RELATION_TYPE).values('dashboard_group_url')
      )  # group_url
      traversal = traversal.by('dashboard_url')  # dashboard_url
      traversal = traversal.by('key')  # uri
  
      traversal = traversal.by(
          __.coalesce(
              __.out('EXECUTED').has('key', TextP.endingWith('_last_successful_execution')).values('timestamp'),
              __.constant('')
          )
      )  # last_successful_run_timestamp
      traversal = traversal.by(
          __.out(DashboardQuery.DASHBOARD_QUERY_RELATION_TYPE).values('name').dedup().fold()
      )  # query_names
      traversal = traversal.by(
          __.out(
              DashboardQuery.DASHBOARD_QUERY_RELATION_TYPE
          ).out(DashboardChart.CHART_RELATION_TYPE).values('name').dedup().fold()
      )  # chart_names
      traversal = traversal.by(__.coalesce(
          __.outE(READ_REVERSE_RELATION_TYPE).values(READ_RELATION_COUNT_PROPERTY),
          __.constant(0)
      ).sum())  # total_usage
      traversal = traversal.by(
          __.out('TAGGED_BY').has('tag_type', 'default').values('keys').dedup().fold()
      )  # tags
      traversal = traversal.by(
          __.out('HAS_BADGE').values('keys').dedup().fold()
      )  # badges
  
      traversal = traversal.order().by(__.select('name'), Order.asc)
  
      dashboards = traversal.toList()
      for dashboard in dashboards:
          dashboard['product'] = dashboard['uri'].split('_')[0]
  
      return dashboards
  
  
  class NeptuneSearchDataExtractor(Extractor):
      """
      Extractor to fetch data required to support search from Neptune's graph database
      """
      QUERY_FUNCTION_CONFIG_KEY = 'query_function'
      QUERY_FUNCTION_KWARGS_CONFIG_KEY = 'query_function_kwargs'
      ENTITY_TYPE_CONFIG_KEY = 'entity_type'
      JOB_PUBLISH_TAG_CONFIG_KEY = 'job_publish_tag'
      MODEL_CLASS_CONFIG_KEY = 'model_class'
  
      DEFAULT_QUERY_BY_ENTITY = {
          'table': _table_search_query,
          'user': _user_search_query,
          'dashboard': _dashboard_search_query
      }
  
      def init(self, conf: ConfigTree) -> None:
          self.conf = conf
          self.entity = conf.get_string(NeptuneSearchDataExtractor.ENTITY_TYPE_CONFIG_KEY, default='table').lower()
  
          if NeptuneSearchDataExtractor.QUERY_FUNCTION_CONFIG_KEY in conf:
              self.query_function = conf.get(NeptuneSearchDataExtractor.QUERY_FUNCTION_CONFIG_KEY)
          else:
              self.query_function = NeptuneSearchDataExtractor.DEFAULT_QUERY_BY_ENTITY[self.entity]
  
          self.job_publish_tag = conf.get_string(NeptuneSearchDataExtractor.JOB_PUBLISH_TAG_CONFIG_KEY, '')
          self.neptune_client = NeptuneSessionClient()
  
          neptune_client_conf = Scoped.get_scoped_conf(conf, self.neptune_client.get_scope())
          self.neptune_client.init(neptune_client_conf)
  
          model_class = conf.get(NeptuneSearchDataExtractor.MODEL_CLASS_CONFIG_KEY, None)
          if model_class:
              module_name, class_name = model_class.rsplit(".", 1)
              mod = importlib.import_module(module_name)
              self.model_class = getattr(mod, class_name)
  
          self._extract_iter: Optional[Any] = None
  
      def close(self) -> None:
          self.neptune_client.close()
  
      def extract(self) -> Optional[Any]:
          if not self._extract_iter:
              self._extract_iter = self._get_extract_iter()
  
          try:
              return next(self._extract_iter)
          except StopIteration:
              return None
  
      def _get_extract_iter(self) -> Any:
          if not hasattr(self, 'results'):
              self.results = self.query_function(self.neptune_client.get_graph(), tag_filter=self.job_publish_tag)
  
          for result in self.results:
              if hasattr(self, 'model_class'):
                  obj = self.model_class(**result)
                  yield obj
              else:
                  yield result
  
      def get_scope(self) -> str:
          return 'extractor.neptune_search_data'
  