content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  """
  The Generic Usage Extractor allows you to populate the "Frequent Users" and "Popular Tables" features within
  Amundsen with the help of the TableColumnUsage class. Because this is a generic usage extractor, you need to create
  a custom log parser to calculate poplarity. There is an example of how to calculate table popularity for Snowflake
  in databuilder/example/scripts/sample_snowflake_table_usage.scala.
  """
  
  import logging
  from typing import (
      Any, Dict, Iterator, Union,
  )
  
  from pyhocon import ConfigFactory, ConfigTree
  
  from databuilder import Scoped
  from databuilder.extractor import sql_alchemy_extractor
  from databuilder.extractor.base_extractor import Extractor
  from databuilder.extractor.sql_alchemy_extractor import SQLAlchemyExtractor
  from databuilder.models.table_column_usage import ColumnReader, TableColumnUsage
  
  LOGGER = logging.getLogger(__name__)
  
  
  class GenericUsageExtractor(Extractor):
      # SELECT statement from table that contains usage data
      SQL_STATEMENT = """
          SELECT
              database,
              schema,
              name,
              user_email,
              read_count
          FROM
              {database}.{schema}.{table}
          {where_clause_suffix};
      """
      # CONFIG KEYS
      WHERE_CLAUSE_SUFFIX_KEY = 'where_clause_suffix'
      POPULARITY_TABLE_DATABASE = 'popularity_table_database'
      POPULARTIY_TABLE_SCHEMA = 'popularity_table_schema'
      POPULARITY_TABLE_NAME = 'popularity_table_name'
      DATABASE_KEY = 'database_key'
  
      DEFAULT_CONFIG = ConfigFactory.from_dict(
          {WHERE_CLAUSE_SUFFIX_KEY: ' ',
           POPULARITY_TABLE_DATABASE: 'PROD',
           POPULARTIY_TABLE_SCHEMA: 'SCHEMA',
           POPULARITY_TABLE_NAME: 'TABLE',
           DATABASE_KEY: 'snowflake'}
      )
  
      def init(self, conf: ConfigTree) -> None:
          conf = conf.with_fallback(GenericUsageExtractor.DEFAULT_CONFIG)
          self._where_clause_suffix = conf.get_string(GenericUsageExtractor.WHERE_CLAUSE_SUFFIX_KEY)
          self._popularity_table_database = conf.get_string(GenericUsageExtractor.POPULARITY_TABLE_DATABASE)
          self._popularity_table_schema = conf.get_string(GenericUsageExtractor.POPULARTIY_TABLE_SCHEMA)
          self._popularity_table_name = conf.get_string(GenericUsageExtractor.POPULARITY_TABLE_NAME)
          self._database_key = conf.get_string(GenericUsageExtractor.DATABASE_KEY)
  
          self.sql_stmt = self.SQL_STATEMENT.format(
              where_clause_suffix=self._where_clause_suffix,
              database=self._popularity_table_database,
              schema=self._popularity_table_schema,
              table=self._popularity_table_name
          )
  
          LOGGER.info("SQL for popularity: {}".format(self.sql_stmt))
  
          self._alchemy_extractor = sql_alchemy_extractor.from_surrounding_config(conf, self.sql_stmt)
          sql_alch_conf = Scoped.get_scoped_conf(conf, self._alchemy_extractor.get_scope()) \
              .with_fallback(ConfigFactory.from_dict({SQLAlchemyExtractor.EXTRACT_SQL: self.sql_stmt}))
          self._alchemy_extractor.init(sql_alch_conf)
          self._extract_iter: Union[None, Iterator] = None
  
      def extract(self) -> Union[TableColumnUsage, None]:
          if not self._extract_iter:
              self._extract_iter = self._get_extract_iter()
          try:
              return next(self._extract_iter)
          except StopIteration:
              return None
  
      def get_scope(self) -> str:
          return "extractor.generic_usage"
  
      def _get_extract_iter(self) -> Iterator[TableColumnUsage]:
          """
          Using raw level iterator, it groups to table and yields TableColumnUsage
          :return:
          """
          for row in self._get_raw_extract_iter():
              col_readers = []
              col_readers.append(ColumnReader(database=self._database_key,
                                              cluster=row["database"],
                                              schema=row["schema"],
                                              table=row["name"],
                                              column="*",
                                              user_email=row["user_email"],
                                              read_count=row["read_count"]))
              yield TableColumnUsage(col_readers=col_readers)
  
      def _get_raw_extract_iter(self) -> Iterator[Dict[str, Any]]:
          """
          Provides iterator of result row from SQLAlchemy extractor
          :return:
          """
          row = self._alchemy_extractor.extract()
          while row:
              yield row
              row = self._alchemy_extractor.extract()
  