content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import json
  import logging
  import os
  from enum import Enum
  from typing import (
      Dict, Iterator, List, Optional, Tuple, Union,
  )
  
  from pyhocon import ConfigTree
  
  from databuilder.extractor.base_extractor import Extractor
  from databuilder.models.badge import Badge, BadgeMetadata
  from databuilder.models.table_lineage import TableLineage
  from databuilder.models.table_metadata import ColumnMetadata, TableMetadata
  from databuilder.models.table_source import TableSource
  
  LOGGER = logging.getLogger(__name__)
  
  
  DBT_CATALOG_REQD_KEYS = ['nodes']
  DBT_MANIFEST_REQD_KEYS = ['nodes', 'child_map']
  DBT_MODEL_TYPE = 'model'
  DBT_MODEL_PREFIX = 'model.'
  DBT_TEST_PREFIX = 'test.'
  
  
  class DBT_TAG_AS(Enum):
      BADGE = 'badge'
      TAG = 'tag'
  
  
  class DBT_MODEL_NAME_KEY(Enum):
      ALIAS = 'alias'
      NAME = 'name'
  
  
  class InvalidDbtInputs(Exception):
      pass
  
  
  class DbtExtractor(Extractor):
      """
      Extracts metadata from the dbt manifest.json and catalog.json files.
      At least one of a manifest or a catalog (or both) must be provided.
      The location of the file or a valid Python dictionary of the content
      can be provided.
  
      Currently the following assets are extracted from these files:
  
      - Tables
      - Columns
      - Definitions
      - Table lineage
      - Tags (converted to Amundsen Badges)
  
      Additional metadagta exists and may be extracted in the future:
  
      - Run / test outcomes
      - Freshness
      - Hooks (as programatic description?)
      - Analysis (as queries for a table??)
      - Table / column level statistics
      - Table comments (as programatic descriptoins)
      """
  
      CATALOG_JSON = "catalog_json"
      MANIFEST_JSON = "manifest_json"
      DATABASE_NAME = 'database_name'
  
      # Dbt Extract Options
      EXTRACT_TABLES = 'extract_tables'
      EXTRACT_DESCRIPTIONS = 'extract_descriptions'
      EXTRACT_TAGS = 'extract_tags'
      EXTRACT_LINEAGE = 'extract_lineage'
      SOURCE_URL = 'source_url'  # Base source code URL for the repo containing dbt workflows
      IMPORT_TAGS_AS = 'import_tags_as'
      SCHEMA_FILTER = 'schema_filter'  # Only extract dbt models from this schema, defaults to all models
      MODEL_NAME_KEY = 'model_name_key'  # Whether to use the "name" or "alias" from dbt as the Amundsen name
  
      # Makes all db, schema, cluster and table names lowercase. This is done so that table metadata from dbt
      # with the default key `Sample://Cluster/Schema/Table` match existing metadata that Amundsen has from
      # the database, which may be `sample://cluster/schema/table`.
      # Most databases that dbt integrates with either use lowercase by default in the information schema
      # or the default Amundsen extractor applies a `lower(...)` function to the result (e.g. snowflake).
      # However, Amundsen does not currently enforce a consistent convention and some databases do support
      # upper and lowercase naming conventions (e.g. Redshift). It may be useful to set this False in the
      # config if the table metadata keys in your database are not all lowercase and to then use a transformer to
      # properly format the string value.
      FORCE_TABLE_KEY_LOWER = 'force_table_key_lower'
  
      def init(self, conf: ConfigTree) -> None:
          self._conf = conf
          self._database_name = conf.get_string(DbtExtractor.DATABASE_NAME)
          self._dbt_manifest = conf.get_string(DbtExtractor.MANIFEST_JSON)
          self._dbt_catalog = conf.get_string(DbtExtractor.CATALOG_JSON)
          # Extract options
          self._extract_tables = conf.get_bool(DbtExtractor.EXTRACT_TABLES, True)
          self._extract_descriptions = conf.get_bool(DbtExtractor.EXTRACT_DESCRIPTIONS, True)
          self._extract_tags = conf.get_bool(DbtExtractor.EXTRACT_TAGS, True)
          self._extract_lineage = conf.get_bool(DbtExtractor.EXTRACT_LINEAGE, True)
          self._source_url = conf.get_string(DbtExtractor.SOURCE_URL, None)
          self._force_table_key_lower = conf.get_bool(DbtExtractor.FORCE_TABLE_KEY_LOWER, True)
          self._dbt_tag_as = DBT_TAG_AS(conf.get_string(DbtExtractor.IMPORT_TAGS_AS, DBT_TAG_AS.BADGE.value))
          self._schema_filter = conf.get_string(DbtExtractor.SCHEMA_FILTER, '')
          self._model_name_key = DBT_MODEL_NAME_KEY(
              conf.get_string(DbtExtractor.MODEL_NAME_KEY, DBT_MODEL_NAME_KEY.NAME.value)).value
          self._clean_inputs()
          self._extract_iter: Union[None, Iterator] = None
  
      def get_scope(self) -> str:
          return "extractor.dbt"
  
      def _validate_catalog(self) -> None:
          # Load the catalog file if needed and run basic validation on the content
          try:
              self._dbt_catalog = json.loads(self._dbt_catalog)
          except Exception:
              try:
                  with open(self._dbt_catalog, 'rb') as f:
                      self._dbt_catalog = json.loads(f.read().lower())
              except Exception as e:
                  raise InvalidDbtInputs(
                      'Invalid content for a dbt catalog was provided. Must be a valid Python '
                      'dictionary or the location of a file. Error received: %s' % e
                  )
          for catalog_key in DBT_CATALOG_REQD_KEYS:
              if catalog_key not in self._dbt_catalog:
                  raise InvalidDbtInputs(
                      "Dbt catalog file must contain keys: %s, found keys: %s"
                      % (DBT_CATALOG_REQD_KEYS, self._dbt_catalog.keys())
                  )
  
      def _validate_manifest(self) -> None:
          # Load the manifest file if needed and run basic validation on the content
          try:
              self._dbt_manifest = json.loads(self._dbt_manifest)
          except Exception:
              try:
                  with open(self._dbt_manifest, 'rb') as f:
                      self._dbt_manifest = json.loads(f.read().lower())
              except Exception as e:
                  raise InvalidDbtInputs(
                      'Invalid content for a dbt manifest was provided. Must be a valid Python '
                      'dictionary or the location of a file. Error received: %s' % e
                  )
          for manifest_key in DBT_MANIFEST_REQD_KEYS:
              if manifest_key not in self._dbt_manifest:
                  raise InvalidDbtInputs(
                      "Dbt manifest file must contain keys: %s, found keys: %s"
                      % (DBT_MANIFEST_REQD_KEYS, self._dbt_manifest.keys())
                  )
  
      def _clean_inputs(self) -> None:
          """
          Validates the dbt input to ensure that at least one of the inputs
          (manifest.json or catalog.json) are provided. Once validated, the
          inputs are sanitized to ensure that the `self._dbt_manifest` and
          `self._dbt_catalog` are valid Python dictionaries.
          """
          if self._database_name is None:
              raise InvalidDbtInputs(
                  'Must provide a database name that corresponds to this dbt catalog and manifest.'
              )
  
          if not self._dbt_manifest or not self._dbt_catalog:
              raise InvalidDbtInputs(
                  'Must provide a dbt manifest file and dbt catalog file.'
              )
  
          self._validate_catalog()
          self._validate_manifest()
  
      def extract(self) -> Union[TableMetadata, None]:
          """
          For every feature table from Feast, a multiple objets are extracted:
  
          1. TableMetadata with feature table description
          2. Programmatic Description of the feature table, containing
             metadata - date of creation and labels
          3. Programmatic Description with Batch Source specification
          4. (if applicable) Programmatic Description with Stream Source
             specification
          """
          if not self._extract_iter:
              self._extract_iter = self._get_extract_iter()
          try:
              return next(self._extract_iter)
          except StopIteration:
              return None
  
      def _default_sanitize(self, s: str) -> str:
          """
          Default function that will be run to convert the value of a string to lowercase.
          """
          if s and self._force_table_key_lower:
              s = s.lower()
          return s
  
      def _get_table_descriptions(self, manifest_content: Dict) -> Tuple[Optional[str], Optional[str]]:
          """
          Gets a description and description source for a table.
          """
          desc, desc_src = None, None
          if self._extract_descriptions:
              desc = manifest_content.get('description')
              desc_src = 'dbt description'
          return desc, desc_src
  
      def _get_table_tags_badges(self, manifest_content: Dict) -> Tuple[Optional[List[str]], Optional[List[str]]]:
          """
          Gets tags or badges for a given table. At most one of these values will not be null.
          """
          tags, tbl_badges = None, None
          if self._extract_tags:
              if self._dbt_tag_as == DBT_TAG_AS.BADGE:
                  tbl_badges = manifest_content.get('tags')
              elif self._dbt_tag_as == DBT_TAG_AS.TAG:
                  tags = manifest_content.get('tags')
          return tags, tbl_badges
  
      def _can_yield_schema(self, schema: str) -> bool:
          """
          Whether or not the schema can be yielded based on the schema filter criteria.
          """
          return (not self._schema_filter) or (self._schema_filter.lower() == schema.lower())
  
      def _get_extract_iter(self) -> Iterator[Union[TableMetadata, BadgeMetadata, TableSource, TableLineage]]:
          """
          Generates the extract iterator for all of the model types created by the dbt files.
          """
          dbt_id_to_table_key = {}
          for tbl_node, manifest_content in self._dbt_manifest['nodes'].items():
  
              if manifest_content['resource_type'] == DBT_MODEL_TYPE and tbl_node in self._dbt_catalog['nodes']:
                  LOGGER.info(
                      'Extracting dbt {}.{}'.format(manifest_content['schema'], manifest_content[self._model_name_key])
                  )
  
                  catalog_content = self._dbt_catalog['nodes'][tbl_node]
  
                  tbl_columns: List[ColumnMetadata] = self._get_column_values(
                      manifest_columns=manifest_content['columns'], catalog_columns=catalog_content['columns']
                  )
  
                  desc, desc_src = self._get_table_descriptions(manifest_content)
                  tags, tbl_badges = self._get_table_tags_badges(manifest_content)
  
                  tbl_metadata = TableMetadata(
                      database=self._default_sanitize(self._database_name),
                      # The dbt "database" is the cluster here
                      cluster=self._default_sanitize(manifest_content['database']),
                      schema=self._default_sanitize(manifest_content['schema']),
                      name=self._default_sanitize(manifest_content[self._model_name_key]),
                      is_view=catalog_content['metadata']['type'] == 'view',
                      columns=tbl_columns,
                      tags=tags,
                      description=desc,
                      description_source=desc_src
                  )
                  # Keep track for Lineage
                  dbt_id_to_table_key[tbl_node] = tbl_metadata._get_table_key()
  
                  # Optionally filter schemas in the output
                  yield_schema = self._can_yield_schema(manifest_content['schema'])
  
                  if self._extract_tables and yield_schema:
                      yield tbl_metadata
  
                  if self._extract_tags and tbl_badges and yield_schema:
                      yield BadgeMetadata(start_label=TableMetadata.TABLE_NODE_LABEL,
                                          start_key=tbl_metadata._get_table_key(),
                                          badges=[Badge(badge, 'table') for badge in tbl_badges])
  
                  if self._source_url and yield_schema:
                      yield TableSource(db_name=tbl_metadata.database,
                                        cluster=tbl_metadata.cluster,
                                        schema=tbl_metadata.schema,
                                        table_name=tbl_metadata.name,
                                        source=os.path.join(self._source_url, manifest_content.get('original_file_path')))
  
          if self._extract_lineage:
              for upstream, downstreams in self._dbt_manifest['child_map'].items():
                  if upstream not in dbt_id_to_table_key:
                      continue
                  valid_downstreams = [
                      dbt_id_to_table_key[k] for k in downstreams
                      if k.startswith(DBT_MODEL_PREFIX) and dbt_id_to_table_key.get(k)
                  ]
                  if valid_downstreams:
                      yield TableLineage(
                          table_key=dbt_id_to_table_key[upstream],
                          downstream_deps=valid_downstreams
                      )
  
      def _get_column_values(self, manifest_columns: Dict, catalog_columns: Dict) -> List[ColumnMetadata]:
          """
          Iterates over the columns in the manifest file and creates a `ColumnMetadata` object
          with the combined information from the manifest file as well as the catalog file.
  
          :params manifest_columns: A dictionary of values from the manifest.json, the keys
              are column names and the values are column metadata
          :params catalog_columns: A dictionary of values from the catalog.json, the keys
              are column names and the values are column metadata
          :returns: A list of `ColumnMetadata` in Amundsen.
          """
          tbl_columns = []
          for catalog_col_name, catalog_col_content in catalog_columns.items():
              manifest_col_content = manifest_columns.get(catalog_col_name, {})
              if catalog_col_content:
                  col_desc = None
                  if self._extract_descriptions:
                      col_desc = manifest_col_content.get('description')
  
                  # Only extract column-level tags IF converting to badges, Amundsen does not have column-level tags
                  badges = None
                  if self._extract_tags and self._dbt_tag_as == DBT_TAG_AS.BADGE:
                      badges = manifest_col_content.get('tags')
  
                  col_metadata = ColumnMetadata(
                      name=self._default_sanitize(catalog_col_content['name']),
                      description=col_desc,
                      col_type=catalog_col_content['type'],
                      sort_order=catalog_col_content['index'],
                      badges=badges
                  )
                  tbl_columns.append(col_metadata)
          return tbl_columns
  