content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import logging
  import time
  from datetime import datetime
  from functools import wraps
  from multiprocessing.pool import ThreadPool
  from typing import (
      Any, Iterator, List, Union,
  )
  
  from pyhocon import ConfigFactory, ConfigTree
  from pytz import UTC
  
  from databuilder import Scoped
  from databuilder.extractor.base_extractor import Extractor
  from databuilder.extractor.sql_alchemy_extractor import SQLAlchemyExtractor
  from databuilder.filesystem.filesystem import FileSystem, is_client_side_error
  from databuilder.models.table_last_updated import TableLastUpdated
  
  LOGGER = logging.getLogger(__name__)
  OLDEST_TIMESTAMP = datetime.fromtimestamp(0, UTC)
  
  
  def fs_error_handler(f: Any) -> Any:
      """
      A Decorator that handles error from FileSystem for HiveTableLastUpdatedExtractor use case
      If it's client side error, it logs in INFO level, and other errors is logged as error level with stacktrace.
      The decorator is intentionally not re-raising exception so that it can isolate the error.
      :param f:
      :return:
      """
  
      @wraps(f)
      def wrapper(*args: Any, **kwargs: Any) -> Any:
          try:
              return f(*args, **kwargs)
          except Exception as e:
              if is_client_side_error(e):
                  LOGGER.info('Invalid metadata. Skipping. args: %s, kwargs: %s. error: %s', args, kwargs, e)
                  return None
              else:
                  LOGGER.exception('Unknown exception while processing args: %s, kwargs: %s', args, kwargs)
                  return None
  
      return wrapper
  
  
  class HiveTableLastUpdatedExtractor(Extractor):
      """
      Uses Hive metastore and underlying storage to figure out last updated timestamp of table.
  
      It turned out that we cannot use table param "last_modified_time", as it provides DDL date not DML date.
      For this reason, we are utilizing underlying file of Hive to fetch latest updated date.
      However, it is not efficient to poke all files in Hive, and it only poke underlying storage for non-partitioned
      table. For partitioned table, it will fetch partition created timestamp, and it's close enough for last updated
      timestamp.
  
      """
      DEFAULT_PARTITION_TABLE_SQL_STATEMENT = """
      SELECT
      DBS.NAME as `schema`,
      TBL_NAME as table_name,
      MAX(PARTITIONS.CREATE_TIME) as last_updated_time
      FROM TBLS
      JOIN DBS ON TBLS.DB_ID = DBS.DB_ID
      JOIN PARTITIONS ON TBLS.TBL_ID = PARTITIONS.TBL_ID
      {where_clause_suffix}
      GROUP BY `schema`, table_name
      ORDER BY `schema`, table_name;
      """
  
      DEFAULT_POSTGRES_PARTITION_TABLE_SQL_STATEMENT = """
      SELECT
      d."NAME" as "schema",
      t."TBL_NAME" as table_name,
      MAX(p."CREATE_TIME") as last_updated_time
      FROM "TBLS" t
      JOIN "DBS" d ON t."DB_ID" = d."DB_ID"
      JOIN "PARTITIONS" p ON t."TBL_ID" = p."TBL_ID"
      {where_clause_suffix}
      GROUP BY "schema", table_name
      ORDER BY "schema", table_name;
      """
  
      DEFAULT_NON_PARTITIONED_TABLE_SQL_STATEMENT = """
      SELECT
      DBS.NAME as `schema`,
      TBL_NAME as table_name,
      SDS.LOCATION as location
      FROM TBLS
      JOIN DBS ON TBLS.DB_ID = DBS.DB_ID
      JOIN SDS ON TBLS.SD_ID = SDS.SD_ID
      {where_clause_suffix}
      ORDER BY `schema`, table_name;
      """
  
      DEFAULT_POSTGRES_NON_PARTITIONED_TABLE_SQL_STATEMENT = """
      SELECT
      d."NAME" as "schema",
      t."TBL_NAME" as table_name,
      s."LOCATION" as location
      FROM "TBLS" t
      JOIN "DBS" d ON t."DB_ID" = d."DB_ID"
      JOIN "SDS" s ON t."SD_ID" = s."SD_ID"
      {where_clause_suffix}
      ORDER BY "schema", table_name;
      """
  
      # Additional where clause for non partitioned table SQL
      DEFAULT_ADDTIONAL_WHERE_CLAUSE = """ NOT EXISTS (SELECT * FROM PARTITIONS WHERE PARTITIONS.TBL_ID = TBLS.TBL_ID)
      AND NOT EXISTS (SELECT * FROM PARTITION_KEYS WHERE PARTITION_KEYS.TBL_ID = TBLS.TBL_ID)
      """
  
      DEFAULT_POSTGRES_ADDTIONAL_WHERE_CLAUSE = """ NOT EXISTS (SELECT * FROM "PARTITIONS" p
      WHERE p."TBL_ID" = t."TBL_ID") AND NOT EXISTS (SELECT * FROM "PARTITION_KEYS" pk WHERE pk."TBL_ID" = t."TBL_ID")
      """
  
      DATABASE = 'hive'
  
      # CONFIG KEYS
      PARTITIONED_TABLE_WHERE_CLAUSE_SUFFIX_KEY = 'partitioned_table_where_clause_suffix'
      NON_PARTITIONED_TABLE_WHERE_CLAUSE_SUFFIX_KEY = 'non_partitioned_table_where_clause_suffix'
      CLUSTER_KEY = 'cluster'
      # number of threads that fetches metadata from FileSystem
      FS_WORKER_POOL_SIZE = 'fs_worker_pool_size'
      FS_WORKER_TIMEOUT_SEC = 'fs_worker_timeout_sec'
      # If number of files that it needs to fetch metadata is larger than this threshold, it will skip the table.
      FILE_CHECK_THRESHOLD = 'file_check_threshold'
  
      DEFAULT_CONFIG = ConfigFactory.from_dict({PARTITIONED_TABLE_WHERE_CLAUSE_SUFFIX_KEY: ' ',
                                                NON_PARTITIONED_TABLE_WHERE_CLAUSE_SUFFIX_KEY: ' ',
                                                CLUSTER_KEY: 'gold',
                                                FS_WORKER_POOL_SIZE: 500,
                                                FS_WORKER_TIMEOUT_SEC: 60,
                                                FILE_CHECK_THRESHOLD: -1})
  
      def init(self, conf: ConfigTree) -> None:
          self._conf = conf.with_fallback(HiveTableLastUpdatedExtractor.DEFAULT_CONFIG)
  
          pool_size = self._conf.get_int(HiveTableLastUpdatedExtractor.FS_WORKER_POOL_SIZE)
          LOGGER.info('Using thread pool size: %s', pool_size)
          self._fs_worker_pool = ThreadPool(processes=pool_size)
          self._fs_worker_timeout = self._conf.get_int(HiveTableLastUpdatedExtractor.FS_WORKER_TIMEOUT_SEC)
          LOGGER.info('Using thread timeout: %s seconds', self._fs_worker_timeout)
  
          self._cluster = self._conf.get_string(HiveTableLastUpdatedExtractor.CLUSTER_KEY)
  
          self._partitioned_table_extractor = self._get_partitioned_table_sql_alchemy_extractor()
          self._non_partitioned_table_extractor = self._get_non_partitioned_table_sql_alchemy_extractor()
          self._fs = self._get_filesystem()
          self._last_updated_filecheck_threshold \
              = self._conf.get_int(HiveTableLastUpdatedExtractor.FILE_CHECK_THRESHOLD)
  
          self._extract_iter: Union[None, Iterator] = None
  
      def _get_partitioned_table_sql_alchemy_extractor(self) -> Extractor:
          """
          Getting an SQLAlchemy extractor that extracts last updated timestamp for partitioned table.
          :return: SQLAlchemyExtractor
          """
  
          sql_stmt = self._choose_default_partitioned_sql_stm().format(
              where_clause_suffix=self._conf.get_string(
                  self.PARTITIONED_TABLE_WHERE_CLAUSE_SUFFIX_KEY, ' '))
  
          LOGGER.info('SQL for partitioned table against Hive metastore: %s', sql_stmt)
  
          sql_alchemy_extractor = SQLAlchemyExtractor()
          sql_alchemy_conf = Scoped.get_scoped_conf(self._conf, sql_alchemy_extractor.get_scope()) \
              .with_fallback(ConfigFactory.from_dict({SQLAlchemyExtractor.EXTRACT_SQL: sql_stmt}))
          sql_alchemy_extractor.init(sql_alchemy_conf)
          return sql_alchemy_extractor
  
      def _choose_default_partitioned_sql_stm(self) -> str:
          conn_string = self._conf.get_string("extractor.sqlalchemy.conn_string")
          if conn_string.startswith('postgres') or conn_string.startswith('postgresql'):
              return self.DEFAULT_POSTGRES_PARTITION_TABLE_SQL_STATEMENT
          else:
              return self.DEFAULT_PARTITION_TABLE_SQL_STATEMENT
  
      def _get_non_partitioned_table_sql_alchemy_extractor(self) -> Extractor:
          """
          Getting an SQLAlchemy extractor that extracts storage location for non-partitioned table for further probing
          last updated timestamp
  
          :return: SQLAlchemyExtractor
          """
  
          default_sql_stmt, default_additional_where_clause = self._choose_default_non_partitioned_sql_stm()
  
          if self.NON_PARTITIONED_TABLE_WHERE_CLAUSE_SUFFIX_KEY in self._conf:
              where_clause_suffix = """
              {}
              AND {}
              """.format(self._conf.get_string(
                  self.NON_PARTITIONED_TABLE_WHERE_CLAUSE_SUFFIX_KEY),
                  default_additional_where_clause)
          else:
              where_clause_suffix = 'WHERE {}'.format(default_additional_where_clause)
  
          sql_stmt = default_sql_stmt.format(
              where_clause_suffix=where_clause_suffix)
  
          LOGGER.info('SQL for non-partitioned table against Hive metastore: %s', sql_stmt)
  
          sql_alchemy_extractor = SQLAlchemyExtractor()
          sql_alchemy_conf = Scoped.get_scoped_conf(self._conf, sql_alchemy_extractor.get_scope()) \
              .with_fallback(ConfigFactory.from_dict({SQLAlchemyExtractor.EXTRACT_SQL: sql_stmt}))
          sql_alchemy_extractor.init(sql_alchemy_conf)
          return sql_alchemy_extractor
  
      def _choose_default_non_partitioned_sql_stm(self) -> List[str]:
          conn_string = self._conf.get_string("extractor.sqlalchemy.conn_string")
          if conn_string.startswith('postgres') or conn_string.startswith('postgresql'):
              return [self.DEFAULT_POSTGRES_NON_PARTITIONED_TABLE_SQL_STATEMENT,
                      self.DEFAULT_POSTGRES_ADDTIONAL_WHERE_CLAUSE]
          else:
              return [self.DEFAULT_NON_PARTITIONED_TABLE_SQL_STATEMENT, self.DEFAULT_ADDTIONAL_WHERE_CLAUSE]
  
      def _get_filesystem(self) -> FileSystem:
          fs = FileSystem()
          fs.init(Scoped.get_scoped_conf(self._conf, fs.get_scope()))
          return fs
  
      def extract(self) -> Union[TableLastUpdated, 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.hive_table_last_updated'
  
      def _get_extract_iter(self) -> Iterator[TableLastUpdated]:
          """
          An iterator that utilizes Generator pattern. First it provides TableLastUpdated objects for partitioned
          table, straight from partitioned_table_extractor (SQLAlchemyExtractor)
  
          Once partitioned table is done, it uses non_partitioned_table_extractor to get storage location of table,
          and probing files under storage location to get max timestamp per table.
          :return:
          """
  
          partitioned_tbl_row = self._partitioned_table_extractor.extract()
          while partitioned_tbl_row:
              yield TableLastUpdated(table_name=partitioned_tbl_row['table_name'],
                                     last_updated_time_epoch=partitioned_tbl_row['last_updated_time'],
                                     schema=partitioned_tbl_row['schema'],
                                     db=HiveTableLastUpdatedExtractor.DATABASE,
                                     cluster=self._cluster)
              partitioned_tbl_row = self._partitioned_table_extractor.extract()
  
          LOGGER.info('Extracting non-partitioned table')
          count = 0
          non_partitioned_tbl_row = self._non_partitioned_table_extractor.extract()
          while non_partitioned_tbl_row:
              count += 1
              if count % 10 == 0:
                  LOGGER.info('Processed %i non-partitioned tables', count)
  
              if not non_partitioned_tbl_row['location']:
                  LOGGER.warning('Skipping as no storage location available. %s', non_partitioned_tbl_row)
                  non_partitioned_tbl_row = self._non_partitioned_table_extractor.extract()
                  continue
  
              start = time.time()
              table_last_updated = self._get_last_updated_datetime_from_filesystem(
                  table=non_partitioned_tbl_row['table_name'],
                  schema=non_partitioned_tbl_row['schema'],
                  storage_location=non_partitioned_tbl_row['location'])
              LOGGER.info('Elapsed: %i seconds', time.time() - start)
  
              if table_last_updated:
                  yield table_last_updated
  
              non_partitioned_tbl_row = self._non_partitioned_table_extractor.extract()
  
      def _get_last_updated_datetime_from_filesystem(self,
                                                     table: str,
                                                     schema: str,
                                                     storage_location: str,
                                                     ) -> Union[TableLastUpdated, None]:
          """
          Fetching metadata within files under storage_location to get latest timestamp.
          (First level only under storage_location)
          Utilizes thread pool to enhance performance. Not using processpool, as it's almost entirely IO bound operation.
  
          :param table:
          :param schema:
          :param storage_location:
          :return:
          """
  
          if LOGGER.isEnabledFor(logging.DEBUG):
              LOGGER.debug(f'Getting last updated datetime for {schema}.{table} in {storage_location}')
  
          last_updated = OLDEST_TIMESTAMP
  
          paths = self._ls(storage_location)
          if not paths:
              LOGGER.info(f'{schema}.{table} does not have any file in path {storage_location}. Skipping')
              return None
  
          LOGGER.info(f'Fetching metadata for {schema}.{table} of {len(paths)} files')
  
          if 0 < self._last_updated_filecheck_threshold < len(paths):
              LOGGER.info(f'Skipping {schema}.{table} due to too many files. '
                          f'{len(paths)} files exist in {storage_location}')
              return None
  
          time_stamp_futures = \
              [self._fs_worker_pool.apply_async(self._get_timestamp, (path, schema, table, storage_location))
               for path in paths]
          for time_stamp_future in time_stamp_futures:
              try:
                  time_stamp = time_stamp_future.get(timeout=self._fs_worker_timeout)
                  if time_stamp:
                      last_updated = max(time_stamp, last_updated)
              except TimeoutError:
                  LOGGER.warning('Timed out on paths %s . Skipping', paths)
  
          if last_updated == OLDEST_TIMESTAMP:
              LOGGER.info(f'No timestamp was derived on {schema}.{table} from location: {storage_location} . Skipping')
              return None
  
          result = TableLastUpdated(table_name=table,
                                    last_updated_time_epoch=int((last_updated - OLDEST_TIMESTAMP).total_seconds()),
                                    schema=schema,
                                    db=HiveTableLastUpdatedExtractor.DATABASE,
                                    cluster=self._cluster)
  
          return result
  
      @fs_error_handler
      def _ls(self, path: str) -> List[str]:
          """
          An wrapper to FileSystem.ls to use fs_error_handler decorator
          :param path:
          :return:
          """
          return self._fs.ls(path)
  
      @fs_error_handler
      def _get_timestamp(self,
                         path: str,
                         schema: str,
                         table: str,
                         storage_location: str,
                         ) -> Union[datetime, None]:
          """
          An wrapper to FileSystem.ls to use fs_error_handler decorator
          :param path:
          :param schema:
          :param table:
          :param storage_location:
          :return:
          """
          if not path:
              LOGGER.info(f'Empty path {path} on {schema}.{table} in storage location {storage_location} . Skipping')
              return None
  
          if not self._fs.is_file(path):
              return None
  
          file_metadata = self._fs.info(path)
          return file_metadata.last_updated
  