content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import copy
  from typing import (
      TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Union,
  )
  
  from amundsen_common.utils.atlas import (
      AtlasCommonParams, AtlasCommonTypes, AtlasTableTypes,
  )
  from amundsen_rds.models import RDSModel
  from amundsen_rds.models.cluster import Cluster as RDSCluster
  from amundsen_rds.models.column import (
      ColumnBadge as RDSColumnBadge, ColumnDescription as RDSColumnDescription, TableColumn as RDSTableColumn,
  )
  from amundsen_rds.models.database import Database as RDSDatabase
  from amundsen_rds.models.schema import Schema as RDSSchema
  from amundsen_rds.models.table import (
      Table as RDSTable, TableDescription as RDSTableDescription,
      TableProgrammaticDescription as RDSTableProgrammaticDescription, TableTag as RDSTableTag,
  )
  from amundsen_rds.models.tag import Tag as RDSTag
  
  from databuilder.models.atlas_entity import AtlasEntity
  from databuilder.models.atlas_relationship import AtlasRelationship
  from databuilder.models.atlas_serializable import AtlasSerializable
  from databuilder.models.badge import Badge, BadgeMetadata
  from databuilder.models.cluster import cluster_constants
  from databuilder.models.description_metadata import (  # noqa: F401
      DESCRIPTION_NODE_LABEL, DESCRIPTION_NODE_LABEL_VAL, DescriptionMetadata,
  )
  from databuilder.models.graph_node import GraphNode
  from databuilder.models.graph_relationship import GraphRelationship
  from databuilder.models.graph_serializable import GraphSerializable
  from databuilder.models.schema import schema_constant
  from databuilder.models.table_serializable import TableSerializable
  
  if TYPE_CHECKING:
      from databuilder.models.type_metadata import TypeMetadata
  
  from databuilder.serializers.atlas_serializer import (
      add_entity_relationship, get_entity_attrs, get_entity_relationships,
  )
  from databuilder.utils.atlas import AtlasRelationshipTypes, AtlasSerializedEntityOperation
  
  
  def _format_as_list(tags: Union[List, str, None]) -> List:
      if tags is None:
          tags = []
      if isinstance(tags, str):
          tags = list(filter(None, tags.split(',')))
      if isinstance(tags, list):
          tags = [tag.lower().strip() for tag in tags]
      return tags
  
  
  class TagMetadata(GraphSerializable, TableSerializable, AtlasSerializable):
      TAG_NODE_LABEL = 'Tag'
      TAG_KEY_FORMAT = '{tag}'
      TAG_TYPE = 'tag_type'
      DEFAULT_TYPE = 'default'
      BADGE_TYPE = 'badge'
      DASHBOARD_TYPE = 'dashboard'
      METRIC_TYPE = 'metric'
  
      TAG_ENTITY_RELATION_TYPE = 'TAG'
      ENTITY_TAG_RELATION_TYPE = 'TAGGED_BY'
  
      def __init__(self,
                   name: str,
                   tag_type: str = 'default',
                   ):
          self._name = name
          self._tag_type = tag_type
          self._nodes = self._create_node_iterator()
          self._relations = self._create_relation_iterator()
          self._records = self._create_record_iterator()
          self._atlas_entity_iterator = self._create_next_atlas_entity()
  
      @staticmethod
      def get_tag_key(name: str) -> str:
          if not name:
              return ''
          return TagMetadata.TAG_KEY_FORMAT.format(tag=name)
  
      def get_node(self) -> GraphNode:
          node = GraphNode(
              key=TagMetadata.get_tag_key(self._name),
              label=TagMetadata.TAG_NODE_LABEL,
              attributes={
                  TagMetadata.TAG_TYPE: self._tag_type
              }
          )
          return node
  
      def get_record(self) -> RDSModel:
          record = RDSTag(
              rk=TagMetadata.get_tag_key(self._name),
              tag_type=self._tag_type
          )
          return record
  
      def create_next_node(self) -> Optional[GraphNode]:
          # return the string representation of the data
          try:
              return next(self._nodes)
          except StopIteration:
              return None
  
      def create_next_relation(self) -> Optional[GraphRelationship]:
          # We don't emit any relations for Tag ingestion
          try:
              return next(self._relations)
          except StopIteration:
              return None
  
      def create_next_record(self) -> Union[RDSModel, None]:
          try:
              return next(self._records)
          except StopIteration:
              return None
  
      def _create_node_iterator(self) -> Iterator[GraphNode]:
          node = self.get_node()
          yield node
  
      def _create_relation_iterator(self) -> Iterator[GraphRelationship]:
          return
          yield
  
      def _create_record_iterator(self) -> Iterator[RDSModel]:
          record = self.get_record()
          yield record
  
      def _create_atlas_glossary_entity(self) -> AtlasEntity:
          attrs_mapping = [
              (AtlasCommonParams.qualified_name, self._name),
              ('glossary', self._tag_type),
              ('term', self._name)
          ]
  
          entity_attrs = get_entity_attrs(attrs_mapping)
  
          entity = AtlasEntity(
              typeName=AtlasCommonTypes.tag,
              operation=AtlasSerializedEntityOperation.CREATE,
              attributes=entity_attrs,
              relationships=None
          )
  
          return entity
  
      def create_atlas_tag_relation(self, table_key: str) -> AtlasRelationship:
          table_relationship = AtlasRelationship(
              relationshipType=AtlasRelationshipTypes.tag,
              entityType1=AtlasCommonTypes.data_set,
              entityQualifiedName1=table_key,
              entityType2=AtlasRelationshipTypes.tag,
              entityQualifiedName2=f'glossary={self._tag_type},term={self._name}',
              attributes={}
          )
  
          return table_relationship
  
      def create_next_atlas_relation(self) -> Union[AtlasRelationship, None]:
          pass
  
      def _create_next_atlas_entity(self) -> Iterator[AtlasEntity]:
          yield self._create_atlas_glossary_entity()
  
      def create_next_atlas_entity(self) -> Union[AtlasEntity, None]:
          try:
              return next(self._atlas_entity_iterator)
          except StopIteration:
              return None
  
  
  class ColumnMetadata:
      COLUMN_NODE_LABEL = 'Column'
      COLUMN_KEY_FORMAT = '{db}://{cluster}.{schema}/{tbl}/{col}'
      COLUMN_NAME = 'name'
      COLUMN_TYPE = 'col_type'
      COLUMN_ORDER = 'sort_order'
      COLUMN_DESCRIPTION = 'description'
      COLUMN_DESCRIPTION_FORMAT = '{db}://{cluster}.{schema}/{tbl}/{col}/{description_id}'
  
      def __init__(self,
                   name: str,
                   description: Union[str, None],
                   col_type: str,
                   sort_order: int,
                   badges: Union[List[str], None] = None,
                   ) -> None:
          """
          TODO: Add stats
          :param name:
          :param description:
          :param col_type:
          :param sort_order:
          :param badges: Optional. Column level badges
          """
          self.name = name
          self.description = DescriptionMetadata.create_description_metadata(source=None,
                                                                             text=description)
          self.type = col_type
          self.sort_order = sort_order
          formatted_badges = _format_as_list(badges)
          self.badges = [Badge(badge, 'column') for badge in formatted_badges]
  
          # The following fields are populated by the ComplexTypeTransformer
          self._column_key: Optional[str] = None
          self._type_metadata: Optional[TypeMetadata] = None
  
      def __repr__(self) -> str:
          return f'ColumnMetadata({self.name!r}, {self.description!r}, {self.type!r}, ' \
                 f'{self.sort_order!r}, {self.badges!r})'
  
      def get_column_key(self) -> Optional[str]:
          return self._column_key
  
      def set_column_key(self, col_key: str) -> None:
          self._column_key = col_key
  
      def get_type_metadata(self) -> Optional['TypeMetadata']:
          return self._type_metadata
  
      def set_type_metadata(self, type_metadata: 'TypeMetadata') -> None:
          self._type_metadata = type_metadata
  
  
  class TableMetadata(GraphSerializable, TableSerializable, AtlasSerializable):
      """
      Table metadata that contains columns. It implements Neo4jCsvSerializable so that it can be serialized to produce
      Table, Column and relation of those along with relationship with table and schema. Additionally, it will create
      Database, Cluster, and Schema with relastionships between those.
      These are being created here as it does not make much sense to have different extraction to produce this. As
      database, cluster, schema would be very repititive with low cardinality, it will perform de-dupe so that publisher
      won't need to publish same nodes, relationships.
  
      This class can be used for both table and view metadata. If it is a View, is_view=True should be passed in.
      """
      TABLE_NODE_LABEL = 'Table'
      TABLE_KEY_FORMAT = '{db}://{cluster}.{schema}/{tbl}'
      TABLE_NAME = 'name'
      IS_VIEW = 'is_view'
  
      TABLE_DESCRIPTION_FORMAT = '{db}://{cluster}.{schema}/{tbl}/{description_id}'
  
      DATABASE_NODE_LABEL = 'Database'
      DATABASE_KEY_FORMAT = 'database://{db}'
      DATABASE_CLUSTER_RELATION_TYPE = cluster_constants.CLUSTER_RELATION_TYPE
      CLUSTER_DATABASE_RELATION_TYPE = cluster_constants.CLUSTER_REVERSE_RELATION_TYPE
  
      CLUSTER_NODE_LABEL = cluster_constants.CLUSTER_NODE_LABEL
      CLUSTER_KEY_FORMAT = '{db}://{cluster}'
      CLUSTER_SCHEMA_RELATION_TYPE = schema_constant.SCHEMA_RELATION_TYPE
      SCHEMA_CLUSTER_RELATION_TYPE = schema_constant.SCHEMA_REVERSE_RELATION_TYPE
  
      SCHEMA_NODE_LABEL = schema_constant.SCHEMA_NODE_LABEL
      SCHEMA_KEY_FORMAT = schema_constant.DATABASE_SCHEMA_KEY_FORMAT
      SCHEMA_TABLE_RELATION_TYPE = 'TABLE'
      TABLE_SCHEMA_RELATION_TYPE = 'TABLE_OF'
  
      TABLE_COL_RELATION_TYPE = 'COLUMN'
      COL_TABLE_RELATION_TYPE = 'COLUMN_OF'
  
      TABLE_TAG_RELATION_TYPE = TagMetadata.ENTITY_TAG_RELATION_TYPE
      TAG_TABLE_RELATION_TYPE = TagMetadata.TAG_ENTITY_RELATION_TYPE
  
      # Only for deduping database, cluster, and schema (table and column will be always processed)
      serialized_nodes_keys: Set[Any] = set()
      serialized_rels_keys: Set[Any] = set()
      serialized_records_keys: Set[Any] = set()
  
      def __init__(self,
                   database: str,
                   cluster: str,
                   schema: str,
                   name: str,
                   description: Union[str, None],
                   columns: Optional[Iterable[ColumnMetadata]] = None,
                   is_view: bool = False,
                   tags: Union[List, str, None] = None,
                   description_source: Union[str, None] = None,
                   **kwargs: Any
                   ) -> None:
          """
          :param database:
          :param cluster:
          :param schema:
          :param name:
          :param description:
          :param columns:
          :param is_view: Indicate whether the table is a view or not
          :param tags:
          :param description_source: Optional. Where the description is coming from. Used to compose unique id.
          :param kwargs: Put additional attributes to the table model if there is any.
          """
          self.database = database
          self.cluster = cluster
          self.schema = schema
          self.name = name
          self.description = DescriptionMetadata.create_description_metadata(text=description, source=description_source)
          self.columns = columns if columns else []
          self.is_view = is_view
          self.attrs: Optional[Dict[str, Any]] = None
  
          self.tags = _format_as_list(tags)
  
          if kwargs:
              self.attrs = copy.deepcopy(kwargs)
  
          self._node_iterator = self._create_next_node()
          self._relation_iterator = self._create_next_relation()
          self._record_iterator = self._create_record_iterator()
          self._atlas_entity_iterator = self._create_next_atlas_entity()
          self._atlas_relation_iterator = self._create_atlas_relation_iterator()
  
      def __repr__(self) -> str:
          return f'TableMetadata({self.database!r}, {self.cluster!r}, {self.schema!r}, {self.name!r} ' \
                 f'{self.description!r}, {self.columns!r}, {self.is_view!r}, {self.tags!r})'
  
      def _get_table_key(self) -> str:
          return TableMetadata.TABLE_KEY_FORMAT.format(db=self.database,
                                                       cluster=self.cluster,
                                                       schema=self.schema,
                                                       tbl=self.name)
  
      def _get_table_description_key(self, description: DescriptionMetadata) -> str:
          return TableMetadata.TABLE_DESCRIPTION_FORMAT.format(db=self.database,
                                                               cluster=self.cluster,
                                                               schema=self.schema,
                                                               tbl=self.name,
                                                               description_id=description.get_description_id())
  
      def _get_database_key(self) -> str:
          return TableMetadata.DATABASE_KEY_FORMAT.format(db=self.database)
  
      def _get_cluster_key(self) -> str:
          return TableMetadata.CLUSTER_KEY_FORMAT.format(db=self.database,
                                                         cluster=self.cluster)
  
      def _get_schema_key(self) -> str:
          return TableMetadata.SCHEMA_KEY_FORMAT.format(db=self.database,
                                                        cluster=self.cluster,
                                                        schema=self.schema)
  
      def _get_col_key(self, col: ColumnMetadata) -> str:
          return ColumnMetadata.COLUMN_KEY_FORMAT.format(db=self.database,
                                                         cluster=self.cluster,
                                                         schema=self.schema,
                                                         tbl=self.name,
                                                         col=col.name,
                                                         badges=col.badges)
  
      def _get_col_description_key(self,
                                   col: ColumnMetadata,
                                   description: DescriptionMetadata) -> str:
          return ColumnMetadata.COLUMN_DESCRIPTION_FORMAT.format(db=self.database,
                                                                 cluster=self.cluster,
                                                                 schema=self.schema,
                                                                 tbl=self.name,
                                                                 col=col.name,
                                                                 description_id=description.get_description_id())
  
      @staticmethod
      def format_tags(tags: Union[List, str, None]) -> List:
          return _format_as_list(tags)
  
      def create_next_node(self) -> Union[GraphNode, None]:
          try:
              return next(self._node_iterator)
          except StopIteration:
              return None
  
      def _create_next_node(self) -> Iterator[GraphNode]:
          yield self._create_table_node()
  
          if self.description:
              node_key = self._get_table_description_key(self.description)
              yield self.description.get_node(node_key)
  
          # Create the table tag nodes
          if self.tags:
              for tag in self.tags:
                  tag_node = TagMetadata(tag).get_node()
                  yield tag_node
  
          for col in self.columns:
              yield from self._create_column_nodes(col)
  
          # Database, cluster, schema
          others = [
              GraphNode(
                  key=self._get_database_key(),
                  label=TableMetadata.DATABASE_NODE_LABEL,
                  attributes={
                      'name': self.database
                  }
              ),
              GraphNode(
                  key=self._get_cluster_key(),
                  label=TableMetadata.CLUSTER_NODE_LABEL,
                  attributes={
                      'name': self.cluster
                  }
              ),
              GraphNode(
                  key=self._get_schema_key(),
                  label=TableMetadata.SCHEMA_NODE_LABEL,
                  attributes={
                      'name': self.schema
                  }
              )
          ]
  
          for node_tuple in others:
              if node_tuple.key not in TableMetadata.serialized_nodes_keys:
                  TableMetadata.serialized_nodes_keys.add(node_tuple.key)
                  yield node_tuple
  
      def _create_table_node(self) -> GraphNode:
          table_attributes = {
              TableMetadata.TABLE_NAME: self.name,
              TableMetadata.IS_VIEW: self.is_view
          }
          if self.attrs:
              for k, v in self.attrs.items():
                  if k not in table_attributes:
                      table_attributes[k] = v
  
          return GraphNode(
              key=self._get_table_key(),
              label=TableMetadata.TABLE_NODE_LABEL,
              attributes=table_attributes
          )
  
      def _create_column_nodes(self, col: ColumnMetadata) -> Iterator[GraphNode]:
          column_node = GraphNode(
              key=self._get_col_key(col),
              label=ColumnMetadata.COLUMN_NODE_LABEL,
              attributes={
                  ColumnMetadata.COLUMN_NAME: col.name,
                  ColumnMetadata.COLUMN_TYPE: col.type,
                  ColumnMetadata.COLUMN_ORDER: col.sort_order
              }
          )
          yield column_node
  
          if col.description:
              node_key = self._get_col_description_key(col, col.description)
              yield col.description.get_node(node_key)
  
          if col.badges:
              col_badge_metadata = BadgeMetadata(
                  start_label=ColumnMetadata.COLUMN_NODE_LABEL,
                  start_key=self._get_col_key(col),
                  badges=col.badges)
              badge_nodes = col_badge_metadata.get_badge_nodes()
              for node in badge_nodes:
                  yield node
  
          type_metadata = col.get_type_metadata()
          if type_metadata:
              yield from type_metadata.create_node_iterator()
  
      def create_next_relation(self) -> Union[GraphRelationship, None]:
          try:
              return next(self._relation_iterator)
          except StopIteration:
              return None
  
      def _create_next_relation(self) -> Iterator[GraphRelationship]:
          schema_table_relationship = GraphRelationship(
              start_key=self._get_schema_key(),
              start_label=TableMetadata.SCHEMA_NODE_LABEL,
              end_key=self._get_table_key(),
              end_label=TableMetadata.TABLE_NODE_LABEL,
              type=TableMetadata.SCHEMA_TABLE_RELATION_TYPE,
              reverse_type=TableMetadata.TABLE_SCHEMA_RELATION_TYPE,
              attributes={}
          )
          yield schema_table_relationship
  
          if self.description:
              yield self.description.get_relation(TableMetadata.TABLE_NODE_LABEL,
                                                  self._get_table_key(),
                                                  self._get_table_description_key(self.description))
  
          if self.tags:
              for tag in self.tags:
                  tag_relationship = GraphRelationship(
                      start_label=TableMetadata.TABLE_NODE_LABEL,
                      start_key=self._get_table_key(),
                      end_label=TagMetadata.TAG_NODE_LABEL,
                      end_key=TagMetadata.get_tag_key(tag),
                      type=TableMetadata.TABLE_TAG_RELATION_TYPE,
                      reverse_type=TableMetadata.TAG_TABLE_RELATION_TYPE,
                      attributes={}
                  )
                  yield tag_relationship
  
          for col in self.columns:
              yield from self._create_column_relations(col)
  
          others = [
              GraphRelationship(
                  start_label=TableMetadata.DATABASE_NODE_LABEL,
                  end_label=TableMetadata.CLUSTER_NODE_LABEL,
                  start_key=self._get_database_key(),
                  end_key=self._get_cluster_key(),
                  type=TableMetadata.DATABASE_CLUSTER_RELATION_TYPE,
                  reverse_type=TableMetadata.CLUSTER_DATABASE_RELATION_TYPE,
                  attributes={}
              ),
              GraphRelationship(
                  start_label=TableMetadata.CLUSTER_NODE_LABEL,
                  end_label=TableMetadata.SCHEMA_NODE_LABEL,
                  start_key=self._get_cluster_key(),
                  end_key=self._get_schema_key(),
                  type=TableMetadata.CLUSTER_SCHEMA_RELATION_TYPE,
                  reverse_type=TableMetadata.SCHEMA_CLUSTER_RELATION_TYPE,
                  attributes={}
              )
          ]
  
          for rel_tuple in others:
              if (rel_tuple.start_key, rel_tuple.end_key, rel_tuple.type) not in TableMetadata.serialized_rels_keys:
                  TableMetadata.serialized_rels_keys.add((rel_tuple.start_key, rel_tuple.end_key, rel_tuple.type))
                  yield rel_tuple
  
      def _create_column_relations(self, col: ColumnMetadata) -> Iterator[GraphRelationship]:
          column_relationship = GraphRelationship(
              start_label=TableMetadata.TABLE_NODE_LABEL,
              start_key=self._get_table_key(),
              end_label=ColumnMetadata.COLUMN_NODE_LABEL,
              end_key=self._get_col_key(col),
              type=TableMetadata.TABLE_COL_RELATION_TYPE,
              reverse_type=TableMetadata.COL_TABLE_RELATION_TYPE,
              attributes={}
          )
          yield column_relationship
  
          if col.description:
              yield col.description.get_relation(
                  ColumnMetadata.COLUMN_NODE_LABEL,
                  self._get_col_key(col),
                  self._get_col_description_key(col, col.description)
              )
  
          if col.badges:
              badge_metadata = BadgeMetadata(start_label=ColumnMetadata.COLUMN_NODE_LABEL,
                                             start_key=self._get_col_key(col),
                                             badges=col.badges)
              badge_relations = badge_metadata.get_badge_relations()
              for relation in badge_relations:
                  yield relation
  
          type_metadata = col.get_type_metadata()
          if type_metadata:
              yield from type_metadata.create_relation_iterator()
  
      def create_next_record(self) -> Union[RDSModel, None]:
          try:
              return next(self._record_iterator)
          except StopIteration:
              return None
  
      def _create_record_iterator(self) -> Iterator[RDSModel]:
          # Database, Cluster, Schema
          others: List[RDSModel] = [
              RDSDatabase(
                  rk=self._get_database_key(),
                  name=self.database
              ),
              RDSCluster(
                  rk=self._get_cluster_key(),
                  name=self.cluster,
                  database_rk=self._get_database_key()
              ),
              RDSSchema(
                  rk=self._get_schema_key(),
                  name=self.schema,
                  cluster_rk=self._get_cluster_key()
              )
          ]
  
          for record in others:
              if record.rk not in TableMetadata.serialized_records_keys:
                  TableMetadata.serialized_records_keys.add(record.rk)
                  yield record
  
          # Table
          yield RDSTable(
              rk=self._get_table_key(),
              name=self.name,
              is_view=self.is_view,
              schema_rk=self._get_schema_key()
          )
  
          # Table description
          if self.description:
              description_record_key = self._get_table_description_key(self.description)
              if self.description.label == DescriptionMetadata.DESCRIPTION_NODE_LABEL:
                  yield RDSTableDescription(
                      rk=description_record_key,
                      description_source=self.description.source,
                      description=self.description.text,
                      table_rk=self._get_table_key()
                  )
              else:
                  yield RDSTableProgrammaticDescription(
                      rk=description_record_key,
                      description_source=self.description.source,
                      description=self.description.text,
                      table_rk=self._get_table_key()
                  )
  
          # Tag
          for tag in self.tags:
              tag_record = TagMetadata(tag).get_record()
              yield tag_record
  
              table_tag_record = RDSTableTag(
                  table_rk=self._get_table_key(),
                  tag_rk=TagMetadata.get_tag_key(tag)
              )
              yield table_tag_record
  
          # Column
          for col in self.columns:
              yield RDSTableColumn(
                  rk=self._get_col_key(col),
                  name=col.name,
                  type=col.type,
                  sort_order=col.sort_order,
                  table_rk=self._get_table_key()
              )
  
              if col.description:
                  description_record_key = self._get_col_description_key(col, col.description)
                  yield RDSColumnDescription(
                      rk=description_record_key,
                      description_source=col.description.source,
                      description=col.description.text,
                      column_rk=self._get_col_key(col)
                  )
  
              if col.badges:
                  badge_metadata = BadgeMetadata(
                      start_label=ColumnMetadata.COLUMN_NODE_LABEL,
                      start_key=self._get_col_key(col),
                      badges=col.badges
                  )
  
                  badge_records = badge_metadata.get_badge_records()
                  for badge_record in badge_records:
                      yield badge_record
  
                      column_badge_record = RDSColumnBadge(
                          column_rk=self._get_col_key(col),
                          badge_rk=badge_record.rk
                      )
                      yield column_badge_record
  
      def _create_atlas_cluster_entity(self) -> AtlasEntity:
          attrs_mapping = [
              (AtlasCommonParams.qualified_name, self._get_cluster_key()),
              ('name', self.cluster),
              ('displayName', self.cluster)
          ]
  
          entity_attrs = get_entity_attrs(attrs_mapping)
  
          entity = AtlasEntity(
              typeName=AtlasCommonTypes.cluster,
              operation=AtlasSerializedEntityOperation.CREATE,
              attributes=entity_attrs,
              relationships=None
          )
  
          return entity
  
      def _create_atlas_database_entity(self) -> AtlasEntity:
          attrs_mapping = [
              (AtlasCommonParams.qualified_name, self._get_database_key()),
              ('name', self.database),
              ('displayName', self.database)
          ]
  
          entity_attrs = get_entity_attrs(attrs_mapping)
  
          relationship_list = []  # type: ignore
  
          add_entity_relationship(
              relationship_list,
              'cluster',
              AtlasCommonTypes.cluster,
              self._get_cluster_key()
          )
  
          entity = AtlasEntity(
              typeName=AtlasTableTypes.database,
              operation=AtlasSerializedEntityOperation.CREATE,
              attributes=entity_attrs,
              relationships=get_entity_relationships(relationship_list)
          )
  
          return entity
  
      def _create_atlas_schema_entity(self) -> AtlasEntity:
          attrs_mapping = [
              (AtlasCommonParams.qualified_name, self._get_schema_key()),
              ('name', self.schema),
              ('displayName', self.schema)
          ]
  
          entity_attrs = get_entity_attrs(attrs_mapping)
  
          relationship_list = []  # type: ignore
  
          add_entity_relationship(
              relationship_list,
              'cluster',
              AtlasCommonTypes.cluster,
              self._get_cluster_key()
          )
  
          entity = AtlasEntity(
              typeName=AtlasTableTypes.schema,
              operation=AtlasSerializedEntityOperation.CREATE,
              attributes=entity_attrs,
              relationships=get_entity_relationships(relationship_list)
          )
  
          return entity
  
      def _create_atlas_table_entity(self) -> AtlasEntity:
          table_type = 'table' if not self.is_view else 'view'
  
          attrs_mapping = [
              (AtlasCommonParams.qualified_name, self._get_table_key()),
              ('name', self.name),
              ('tableType', table_type),
              ('description', self.description.text if self.description else ''),
              ('displayName', self.name)
          ]
  
          entity_attrs = get_entity_attrs(attrs_mapping)
  
          relationship_list = []  # type: ignore
  
          add_entity_relationship(
              relationship_list,
              'amundsen_schema',
              AtlasTableTypes.schema,
              self._get_schema_key()
          )
  
          entity = AtlasEntity(
              typeName=AtlasTableTypes.table,
              operation=AtlasSerializedEntityOperation.CREATE,
              attributes=entity_attrs,
              relationships=get_entity_relationships(relationship_list)
          )
  
          return entity
  
      def _create_atlas_column_entity(self, column_metadata: ColumnMetadata) -> AtlasEntity:
          qualified_name = column_metadata.COLUMN_KEY_FORMAT.format(db=self.database,
                                                                    cluster=self.cluster,
                                                                    schema=self.schema,
                                                                    tbl=self.name,
                                                                    col=column_metadata.name)
          attrs_mapping = [
              (AtlasCommonParams.qualified_name, qualified_name),
              ('name', column_metadata.name or ''),
              ('description', column_metadata.description.text if column_metadata.description else ''),
              ('type', column_metadata.type),
              ('position', column_metadata.sort_order),
              ('displayName', column_metadata.name or '')
          ]
  
          entity_attrs = get_entity_attrs(attrs_mapping)
  
          relationship_list = []  # type: ignore
  
          add_entity_relationship(
              relationship_list,
              'table',
              AtlasTableTypes.table,
              self._get_table_key()
          )
  
          entity = AtlasEntity(
              typeName=AtlasTableTypes.column,
              operation=AtlasSerializedEntityOperation.CREATE,
              attributes=entity_attrs,
              relationships=get_entity_relationships(relationship_list)
          )
  
          return entity
  
      def _create_atlas_relation_iterator(self) -> Iterator[AtlasRelationship]:
          for tag in self.tags:
              tag_relation = TagMetadata(tag).create_atlas_tag_relation(self._get_table_key())
              yield tag_relation
  
      def create_next_atlas_relation(self) -> Union[AtlasRelationship, None]:
          try:
              return next(self._atlas_relation_iterator)
          except StopIteration:
              return None
  
      def _create_next_atlas_entity(self) -> Iterator[AtlasEntity]:
          yield self._create_atlas_cluster_entity()
          yield self._create_atlas_database_entity()
          yield self._create_atlas_schema_entity()
          yield self._create_atlas_table_entity()
  
          for col in self.columns:
              yield self._create_atlas_column_entity(col)
  
          if self.tags:
              for tag in self.tags:
                  tag_entity = TagMetadata(tag).create_next_atlas_entity()
                  if tag_entity:
                      yield tag_entity
  
      def create_next_atlas_entity(self) -> Union[AtlasEntity, None]:
          try:
              return next(self._atlas_entity_iterator)
          except StopIteration:
              return None
  