content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import json
  import logging
  import re
  from typing import Dict
  
  from apache_atlas.client.base_client import AtlasClient
  from apache_atlas.exceptions import AtlasServiceException
  from apache_atlas.model.misc import SearchFilter
  from apache_atlas.model.typedef import AtlasTypesDef
  from requests import Timeout
  
  from databuilder.types.atlas.types_def import (
      application_schema, bookmark_schema, cluster_schema, column_schema, column_table_relation, dashboard_chart_schema,
      dashboard_execution_schema, dashboard_group_schema, dashboard_query_schema, dashboard_schema, data_owner_schema,
      database_cluster_relation, database_schema, hive_table_partition, lineage_schema, reader_referenceable_relation,
      reader_schema, report_schema, schema_cluster_relation, schema_schema, source_schema, table_partition_schema,
      table_schema, table_schema_relation, table_source_relation, user_reader_relation, user_schema,
  )
  
  LOGGER = logging.getLogger(__name__)
  
  
  # noinspection PyMethodMayBeStatic
  class AtlasEntityInitializer:
      def __init__(self, client: AtlasClient) -> None:
          self.driver = client
  
      def assign_subtypes(self, regex: str, super_type: str) -> None:
          LOGGER.info(f'\nAssigning {super_type} entity to all the subtypes entity definitions with postfix ')
          entities_to_update = []
          entity_defs = self.driver.typedef.get_all_typedefs(search_filter=SearchFilter()).get('entityDefs', [])
  
          for e in entity_defs:
              if re.compile(regex).match(e.name) is not None:
                  LOGGER.info(f'Assigning {e.name} as a subtype of {super_type}')
                  e["superTypes"].append(super_type)
                  entities_to_update.append(e)
  
          typedef_dict = {
              "entityDefs": entities_to_update
          }
  
          self.driver.typedef.update_atlas_typedefs(AtlasTypesDef(attrs=typedef_dict))
          LOGGER.info(f'Assignment of "{super_type}" Entity to existing "{regex}" entities Completed.\n')
  
      def create_or_update(self, typedef_dict: Dict, info: str, attempt: int = 1) -> None:
          try:
              LOGGER.info(f"Trying to create {info} Entity")
              self.driver.typedef.create_atlas_typedefs(AtlasTypesDef(attrs=typedef_dict))
          except AtlasServiceException:
              LOGGER.info(f"Already Exists, updating {info} Entity")
              try:
                  self.driver.typedef.update_atlas_typedefs(AtlasTypesDef(attrs=typedef_dict))
              except Exception:
                  # This is a corner case, for Atlas Sample Data
                  LOGGER.warn(f"Error updating {info} Entity.", exc_info=True)
  
          except Timeout:
              # Sometimes on local atlas instance you do get ReadTimeout a lot.
              # This will try to apply definition 3 times and then cancel
              if attempt < 4:
                  LOGGER.info("ReadTimeout - Another Try.")
                  self.create_or_update(typedef_dict, info, attempt + 1)
              else:
                  LOGGER.info(f"ReadTimeout Exception - Cancelling Operation: {attempt}", exc_info=True)
          except Exception:
              LOGGER.info(f"Error creating/updating {info} Entity Definition", exc_info=True)
          finally:
              LOGGER.info(f"Applied {info} Entity Definition")
  
      def get_schema_dict(self, schema: str) -> Dict:
          return json.loads(schema)
  
      def create_table_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(table_schema), "Table")
  
      def create_column_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(column_schema), "Column")
  
      def create_column_table_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(column_table_relation), "Column <-> Table")
  
      def create_cluster_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(cluster_schema), "Cluster")
  
      def create_database_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(database_schema), "Database")
  
      def create_database_cluster_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(database_cluster_relation), "Database <-> Cluster")
  
      def create_schema_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(schema_schema), "Schema")
  
      def create_schema_cluster_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(schema_cluster_relation), "Schema <-> Database")
  
      def create_table_schema_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(table_schema_relation), "Table <-> Schema")
  
      def create_user_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(user_schema), "User")
  
      def create_reader_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(reader_schema), "Reader")
  
      def create_bookmark_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(bookmark_schema), "Bookmark")
  
      def create_report_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(report_schema), "Report")
  
      def create_user_reader_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(user_reader_relation), "User <-> Reader")
  
      def create_reader_referenceable_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(reader_referenceable_relation), "Reader <-> Referenceable")
  
      def create_table_partition_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(table_partition_schema), "Partition")
  
      def create_hive_table_partition(self) -> None:
          self.create_or_update(self.get_schema_dict(hive_table_partition), "Hive Table Partition")
  
      def create_data_owner_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(data_owner_schema), "Data Owner Relation")
  
      def create_application_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(application_schema), "Application")
  
      def create_source_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(source_schema), "Source")
  
      def create_table_source_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(table_source_relation), "Table <-> Source")
  
      def create_lineage_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(lineage_schema), "LineageProcess")
  
      def create_dashboard_group_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(dashboard_group_schema), "Dashboard Group")
  
      def create_dashboard_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(dashboard_schema), "Dashboard")
  
      def create_dashboard_chart_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(dashboard_chart_schema), "Dashboard Chart")
  
      def create_dashboard_query_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(dashboard_query_schema), "Dashboard Query")
  
      def create_dashboard_execution_schema(self) -> None:
          self.create_or_update(self.get_schema_dict(dashboard_execution_schema), "Dashboard Execution")
  
      def create_dashboard_cluster_relation(self) -> None:
          self.create_or_update(self.get_schema_dict(database_cluster_relation), "Dashboard <-> Cluster")
  
      def create_required_entities(self, fix_existing_data: bool = False) -> None:
          """
          IMPORTANT: The order of the entity definition matters.
          Please keep this order.
          :return: Creates or Updates the entity definition in Apache Atlas
          """
          self.create_cluster_schema()
          self.create_column_schema()
          self.create_reader_schema()
          self.create_user_schema()
          self.create_bookmark_schema()
          self.create_report_schema()
          self.create_database_schema()
          self.create_database_cluster_relation()
          self.create_schema_schema()
          self.create_schema_cluster_relation()
          self.create_table_schema()
          self.create_column_table_relation()
          self.create_table_schema_relation()
          self.create_source_schema()
          self.create_table_source_relation()
          self.create_application_schema()
          self.create_lineage_schema()
          self.assign_subtypes(regex="(.*)_table$", super_type="Table")
          self.assign_subtypes(regex="(.*)_column$", super_type="Column")
          self.create_user_reader_relation()
          self.create_reader_referenceable_relation()
          self.create_table_partition_schema()
          self.create_hive_table_partition()
          self.create_data_owner_relation()
          self.create_dashboard_group_schema()
          self.create_dashboard_schema()
          self.create_dashboard_query_schema()
          self.create_dashboard_chart_schema()
          self.create_dashboard_execution_schema()
          self.create_dashboard_cluster_relation()
  