content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  import abc
  import re
  from typing import Any, Dict, Optional, Set
  
  
  class AtlasStatus:
      ACTIVE = "ACTIVE"
      DELETED = "DELETED"
  
  
  class AtlasCommonParams:
      qualified_name = 'qualifiedName'
      guid = 'guid'
      attributes = 'attributes'
      relationships = 'relationshipAttributes'
      uri = 'entityUri'
      type_name = 'typeName'
      unique_attributes = 'uniqueAttributes'
      created_timestamp = 'createdTimestamp'
      last_modified_timestamp = 'lastModifiedTimestamp'
  
  
  class AtlasCommonTypes:
      bookmark = 'Bookmark'
      user = 'User'
      reader = 'Reader'
      cluster = 'Cluster'
      application = 'Application'
      data_set = 'DataSet'
  
      # These are just `virtual` types which do not actually exist in Atlas.
      # We use those constant values to distinguish Atlas Python Client methods which should be used for populating
      # such data.
      # Tags are published using Glossary API, badges using Classification API. Other entities are published using regular
      # Entity API.
      tag = 'Tag'
      badge = 'Badge'
      resource_report = 'Report'
  
  
  class AtlasTableTypes:
      table = 'Table'
      column = 'Column'
      database = 'Database'
      schema = 'Schema'
      source = 'Source'
      watermark = 'TablePartition'
      process = 'LineageProcess'
  
  
  class AtlasDashboardTypes:
      metadata = 'Dashboard'
      group = 'DashboardGroup'
      query = 'DashboardQuery'
      chart = 'DashboardChart'
      execution = 'DashboardExecution'
  
  
  class AtlasKey(abc.ABC):
      """
      Class for unification of entity keys between Atlas and Amundsen ecosystems.
  
      Since Atlas can be populated both by tools from 'Atlas world' (like Apache Atlas Hive hook/bridge) and Amundsen
      Databuilder (and each of the approach has a different way to render unique identifiers) we need such class
      to serve as unification layer.
      """
  
      def __init__(self, raw_id: str, database: Optional[str] = None) -> None:
          self._raw_identifier = raw_id
          self._database = database
  
      @property
      def is_qualified_name(self) -> bool:
          """
          Property assessing whether raw_id is qualified name.
  
          :returns: -
          """
          if self.atlas_qualified_name_regex.match(self._raw_identifier):
              return True
          else:
              return False
  
      @property
      def is_amundsen_key(self) -> bool:
          """
          Property assessing whether raw_id is amundsen key.
  
          :returns: -
          """
          if self.amundsen_key_regex.match(self._raw_identifier):
              return True
          else:
              return False
  
      def get_details(self) -> Dict[str, str]:
          """
          Collect as many details from key (either qn or amundsen key)
  
          :returns: dictionary of entity properties derived from key
          """
          if self.is_qualified_name:
              return self._get_details_from_qualified_name()
          elif self.is_amundsen_key:
              return self._get_details_from_key()
          else:
              raise ValueError(f'Value is neither valid qualified name nor amundsen key: {self._raw_identifier}')
  
      def _get_details(self, pattern: Any) -> Dict[str, str]:
          """
          Helper function collecting data from regex match
  
          :returns: dictionary of matched regex groups with their values
          """
          try:
              result = pattern.match(self._raw_identifier).groupdict()
  
              return result
          except KeyError:
              raise KeyError
  
      def _get_details_from_qualified_name(self) -> Dict[str, str]:
          """
          Collect as many details from qualified name
  
          :returns: dictionary of entity properties derived from qualified name
          """
          try:
              return self._get_details(self.atlas_qualified_name_regex)
          except KeyError:
              raise ValueError(f'This is not valid qualified name: {self._raw_identifier}')
  
      def _get_details_from_key(self) -> Dict[str, str]:
          """
          Collect as many details from amundsen key
  
          :returns: dictionary of entity properties derived from amundsen key
          """
          try:
              return self._get_details(self.amundsen_key_regex)
          except KeyError:
              raise ValueError(f'This is not valid qualified name: {self._raw_identifier}')
  
      @property
      @abc.abstractmethod
      def atlas_qualified_name_regex(self) -> Any:
          """
          Regex for validating qualified name (and collecting details from qn parts)
  
          :returns: -
          """
          pass
  
      @property
      @abc.abstractmethod
      def amundsen_key_regex(self) -> Any:
          """
          Regex for validating amundsen key (and collecting details from key parts)
  
          :returns: -
          """
          pass
  
      @property
      @abc.abstractmethod
      def qualified_name(self) -> str:
          """
          Properly formatted qualified name
  
          :returns: -
          """
          pass
  
      @property
      @abc.abstractmethod
      def amundsen_key(self) -> str:
          """
          Properly formetted amundsen key
  
          :returns: -
          """
          pass
  
      @property
      def native_atlas_entity_types(self) -> Set[str]:
          """
          Atlas can be populated using two approaches:
          1. Using Atlas-provided bridge/hook tools to ingest data in push manner (like Atlas Hive Hook)
          2. Using Amundsen-provided databuilder framework in pull manner
  
          Since Atlas-provided tools follow different approach for rendering qualified name than databuilder does,
          to provide compatibility for both approaches we need to act differently depending whether the table entity
          was loaded by Atlas-provided or Amundsen-provided tools. We distinguish them by entity type - in Atlas the
          naming convention assumes '_table' suffix in entity name while Amundsen does not have such suffix.
  
          If the entity_type (database in Amundsen lingo) is one of the values from this property, we treat it like
          it was provided by Atlas and follow Atlas qualified name convention.
  
          If the opposite is true - we treat it like it was provided by Amundsen Databuilder, use generic entity types
          and follow Amundsen key name convention.
          """
          return {'hive_table'}
  
      @property
      def entity_type(self) -> str:
          if self.is_qualified_name:
              return self._database or ''
          else:
              return self.get_details()['database'] \
                  if self.get_details()['database'] in self.native_atlas_entity_types else 'Table'
  
  
  class AtlasTableKey(AtlasKey):
      @property
      def atlas_qualified_name_regex(self) -> Any:
          return re.compile(r'^(?P<schema>.*?)\.(?P<table>.*)@(?P<cluster>.*?)$', re.X)
  
      @property
      def amundsen_key_regex(self) -> Any:
          return re.compile(r'^(?P<database>.*?)://(?P<cluster>.*)\.(?P<schema>.*?)\/(?P<table>.*?)$', re.X)
  
      @property
      def qualified_name(self) -> str:
          if not self.is_qualified_name and self.get_details()['database'] in self.native_atlas_entity_types:
              spec = self._get_details_from_key()
  
              schema = spec['schema']
              table = spec['table']
              cluster = spec['cluster']
  
              return f'{schema}.{table}@{cluster}'
          else:
              return self._raw_identifier
  
      @property
      def amundsen_key(self) -> str:
          if self.is_qualified_name:
              spec = self._get_details_from_qualified_name()
  
              schema = spec['schema']
              table = spec['table']
              cluster = spec['cluster']
  
              return f'{self._database}://{cluster}.{schema}/{table}'
          elif self.is_amundsen_key:
              return self._raw_identifier
          else:
              raise ValueError(f'Value is neither qualified name nor amundsen key: {self._raw_identifier}')
  
  
  class AtlasColumnKey(AtlasKey):
      @property
      def atlas_qualified_name_regex(self) -> Any:
          return re.compile(r'^(?P<schema>.*?)\.(?P<table>.*?)\.(?P<column>.*?)@(?P<cluster>.*?)$', re.X)
  
      @property
      def amundsen_key_regex(self) -> Any:
          return re.compile(r'^(?P<database>.*?)://(?P<cluster>.*)\.(?P<schema>.*?)\/(?P<table>.*?)\/(?P<column>.*)$',
                            re.X)
  
      @property
      def qualified_name(self) -> str:
          if self.is_amundsen_key:
              spec = self._get_details_from_key()
  
              schema = spec['schema']
              table = spec['table']
              cluster = spec['cluster']
              column = spec['column']
  
              return f'{schema}.{table}.{column}@{cluster}'
          elif self.is_qualified_name:
              return self._raw_identifier
          else:
              raise ValueError(f'Value is neither qualified name nor amundsen key: {self._raw_identifier}')
  
      @property
      def amundsen_key(self) -> str:
          if self.is_qualified_name:
              spec = self._get_details_from_qualified_name()
  
              schema = spec['schema']
              table = spec['table']
              cluster = spec['cluster']
              column = spec['column']
  
              source = self._database.replace('column', 'table') if self._database else ''
  
              return f'{source}://{cluster}.{schema}/{table}/{column}'
          elif self.is_amundsen_key:
              return self._raw_identifier
          else:
              raise ValueError(f'Value is neither qualified name nor amundsen key: {self._raw_identifier}')
  