content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import logging
  import re
  from collections import namedtuple
  from datetime import (
      datetime, timedelta, timezone,
  )
  from time import sleep
  from typing import (
      Any, Dict, Iterator, List, Optional, Tuple,
  )
  
  from pyhocon import ConfigTree
  
  from databuilder.extractor.base_bigquery_extractor import BaseBigQueryExtractor
  
  TableColumnUsageTuple = namedtuple('TableColumnUsageTuple', ['database', 'cluster', 'schema',
                                                               'table', 'column', 'email'])
  
  LOGGER = logging.getLogger(__name__)
  
  
  class BigQueryTableUsageExtractor(BaseBigQueryExtractor):
      """
      An aggregate extractor for bigquery table usage. This class takes the data from
      the stackdriver logging API by filtering on timestamp, bigquery_resource and looking
      for referencedTables in the response.
      """
      TIMESTAMP_KEY = 'timestamp'
      _DEFAULT_SCOPES = ['https://www.googleapis.com/auth/cloud-platform']
      EMAIL_PATTERN = 'email_pattern'
      DELAY_TIME = 'delay_time'
      TABLE_DECORATORS = ['$', '@']
      COUNT_READS_ONLY_FROM_PROJECT_ID_KEY = 'count_reads_only_from_project_id_key'
  
      def init(self, conf: ConfigTree) -> None:
          BaseBigQueryExtractor.init(self, conf)
          self.timestamp = conf.get_string(
              BigQueryTableUsageExtractor.TIMESTAMP_KEY,
              (datetime.now(timezone.utc) - timedelta(days=1)).strftime(BigQueryTableUsageExtractor.DATE_TIME_FORMAT))
  
          self.email_pattern = conf.get_string(BigQueryTableUsageExtractor.EMAIL_PATTERN, None)
          self.delay_time = conf.get_int(BigQueryTableUsageExtractor.DELAY_TIME, 100)
          self.table_usage_counts: Dict[TableColumnUsageTuple, int] = {}
          # GCP console allows running queries using tables from a project different from the one the extractor is
          # used for; only usage metadata of referenced tables present in the given project_id_key for the
          # extractor is taken into account and usage metadata of referenced tables from other projects
          # is ignored by "default".
          self.count_reads_only_from_same_project = conf.get_bool(
              BigQueryTableUsageExtractor.COUNT_READS_ONLY_FROM_PROJECT_ID_KEY, True)
          self._count_usage()
          self.iter = iter(self.table_usage_counts)
  
      def _count_usage(self) -> None:  # noqa: C901
          count = 0
          for entry in self._retrieve_records():
              count += 1
              if count % self.pagesize == 0:
                  LOGGER.info(f'Aggregated {count} records')
  
              if entry is None:
                  continue
  
              try:
                  job = entry['protoPayload']['serviceData']['jobCompletedEvent']['job']
              except Exception:
                  continue
              if job['jobStatus']['state'] != 'DONE':
                  # This job seems not to have finished yet, so we ignore it.
                  continue
              if len(job['jobStatus'].get('error', {})) > 0:
                  # This job has errors, so we ignore it
                  continue
  
              email = entry['protoPayload']['authenticationInfo']['principalEmail']
              # Query results can be cached and if the source tables remain untouched,
              # bigquery will return it from a 24 hour cache result instead. In that
              # case, referencedTables has been observed to be empty:
              # https://cloud.google.com/logging/docs/reference/audit/bigquery/rest/Shared.Types/AuditData#JobStatistics
  
              refTables = job['jobStatistics'].get('referencedTables', None)
              if refTables:
                  if 'totalTablesProcessed' in job['jobStatistics']:
                      self._create_records(
                          refTables,
                          job['jobStatistics']['totalTablesProcessed'], email,
                          job['jobName']['jobId'])
  
              refViews = job['jobStatistics'].get('referencedViews', None)
              if refViews:
                  if 'totalViewsProcessed' in job['jobStatistics']:
                      self._create_records(
                          refViews, job['jobStatistics']['totalViewsProcessed'],
                          email, job['jobName']['jobId'])
  
      def _create_records(self, refResources: List[dict], resourcesProcessed: int, email: str,
                          jobId: str) -> None:
          # if email filter is provided, only the email matched with filter will be recorded.
          if self.email_pattern:
              if not re.match(self.email_pattern, email):
                  # the usage account not match email pattern
                  return
  
          if len(refResources) != resourcesProcessed:
              LOGGER.warning(f'The number of tables listed in job {jobId} is not consistent')
              return
  
          for refResource in refResources:
              tableId = refResource.get('tableId')
              datasetId = refResource.get('datasetId')
  
              if not datasetId or not tableId:
                  # handling case when the referenced table is an external table
                  # Which doesn't have a datasetId
                  continue
  
              if self._is_anonymous_dataset(datasetId) or self._is_wildcard_table(tableId):
                  continue
  
              tableId = self._remove_table_decorators(tableId)
  
              if self._is_sharded_table(tableId):
                  # Use the prefix of the sharded table as tableId
                  tableId = tableId[:-len(self._get_sharded_table_suffix(tableId))]
  
              if refResource['projectId'] != self.project_id and self.count_reads_only_from_same_project:
                  LOGGER.debug(
                      f'Not counting usage for {refResource} since {tableId} '
                      f'is not present in {self.project_id} '
                      f'and {BigQueryTableUsageExtractor.COUNT_READS_ONLY_FROM_PROJECT_ID_KEY} is True')
                  continue
              else:
                  key = TableColumnUsageTuple(database='bigquery',
                                              cluster=refResource['projectId'],
                                              schema=datasetId,
                                              table=tableId,
                                              column='*',
                                              email=email)
  
              new_count = self.table_usage_counts.get(key, 0) + 1
              self.table_usage_counts[key] = new_count
  
      def _retrieve_records(self) -> Iterator[Optional[Dict]]:
          """
          Extracts bigquery log data by looking at the principalEmail in the authenticationInfo block and
          referencedTables in the jobStatistics and filters out log entries of metadata queries.
          :return: Provides a record or None if no more to extract
          """
          body = {
              'resourceNames': [f'projects/{self.project_id}'],
              'pageSize': self.pagesize,
              'filter': 'protoPayload.methodName="jobservice.jobcompleted" AND '
                        'resource.type="bigquery_resource" AND '
                        'NOT protoPayload.serviceData.jobCompletedEvent.job.jobConfiguration.query.query:('
                        'INFORMATION_SCHEMA OR __TABLES__) AND '
                        f'timestamp >= "{self.timestamp}" AND timestamp < "{self.cutoff_time}"'
          }
          for page in self._page_over_results(body):
              for entry in page['entries']:
                  yield entry
  
      def extract(self) -> Optional[Tuple[Any, int]]:
          try:
              key = next(self.iter)
              return key, self.table_usage_counts[key]
          except StopIteration:
              return None
  
      def _page_over_results(self, body: Dict) -> Iterator[Dict]:
          response = self.logging_service.entries().list(body=body).execute(
              num_retries=BigQueryTableUsageExtractor.NUM_RETRIES)
          while response:
              if 'entries' in response:
                  yield response
  
              try:
                  if 'nextPageToken' in response:
                      body['pageToken'] = response['nextPageToken']
                      response = self.logging_service.entries().list(body=body).execute(
                          num_retries=BigQueryTableUsageExtractor.NUM_RETRIES)
                  else:
                      response = None
              except Exception:
                  # Add a delay when BQ quota exceeds limitation
                  sleep(self.delay_time)
  
      def _remove_table_decorators(self, tableId: str) -> Optional[str]:
          for decorator in BigQueryTableUsageExtractor.TABLE_DECORATORS:
              tableId = tableId.split(decorator)[0]
          return tableId
  
      def _is_anonymous_dataset(self, datasetId: str) -> bool:
          # temporary/cached results tables are stored in anonymous datasets that have names starting with '_'
          return datasetId.startswith('_')
  
      def _is_wildcard_table(self, tableId: str) -> bool:
          return '*' in tableId
  
      def get_scope(self) -> str:
          return 'extractor.bigquery_table_usage'
  