gcpdiag.queries.bigquery

66def get_project_policy(context: models.Context):
67  """Fetches the IAM policy object for a project."""
68  root_logger = logging.getLogger()
69  original_level = root_logger.level
70
71  try:
72    root_logger.setLevel(logging.ERROR)
73    policy = iam.get_project_policy(context, raise_error_if_fails=False)
74    return policy
75  except utils.GcpApiError:
76    return None
77  finally:
78    root_logger.setLevel(original_level)

 81def get_organization_policy(context: models.Context, organization_id: str):
 82  """Fetches the IAM policy object for an organization."""
 83  root_logger = logging.getLogger()
 84  original_level = root_logger.level
 85
 86  try:
 87    root_logger.setLevel(logging.ERROR)
 88    policy = iam.get_organization_policy(context,
 89                                         organization_id,
 90                                         raise_error_if_fails=False)
 91    return policy
 92  except utils.GcpApiError as err:
 93    if 'doesn\'t have access to' in err.message.lower(
 94    ) or 'denied on resource' in err.message.lower():
 95      op.info(
 96          'User does not have access to the organization policy. Investigation'
 97          ' completeness and accuracy might depend on the presence of'
 98          ' organization level permissions.')
 99    return None
100  finally:
101    root_logger.setLevel(original_level)

104def check_permissions_for_principal(
105    policy: PolicyObject, principal: str,
106    permissions_to_check: Set[str]) -> Dict[str, bool]:
107  """Uses a policy object to check a set of permissions for a principal.
108
109  Returns a dictionary mapping each permission to a boolean indicating its
110  presence.
111  """
112  return {
113      permission: policy.has_permission(principal, permission)
114      for permission in permissions_to_check
115  }

118def get_missing_permissions(required_permissions: Set[str],
119                            actual_permissions: Dict[str, bool]) -> Set[str]:
120  """Compares a set of required permissions against a dictionary of actual
121
122  permissions and returns the set of missing ones.
123  """
124  return {
125      perm for perm in required_permissions if not actual_permissions.get(perm)
126  }

129class BigQueryTable:
130  """Represents a BigQuery Table object."""
131
132  project_id: str
133  dataset_id: str
134  table_id: str
135
136  def __init__(self, project_id: str, dataset_id: str, table_id: str):
137    self.project_id = project_id
138    self.dataset_id = dataset_id
139    self.table_id = table_id
140
141  @property
142  def table_identifier(self) -> str:
143    return f'{self.project_id}:{self.dataset_id}.{self.table_id}'

146class BigQueryRoutine:
147  """Represents a BigQuery Routine object."""
148
149  project_id: str
150  dataset_id: str
151  routine_id: str
152
153  def __init__(self, project_id: str, dataset_id: str, routine_id: str):
154    self.project_id = project_id
155    self.dataset_id = dataset_id
156    self.routine_id = routine_id
157
158  @property
159  def routine_identifier(self) -> str:
160    return f'{self.project_id}:{self.dataset_id}.{self.routine_id}'

163class BigQueryJob(models.Resource):
164  """Represents a BigQuery Job object."""
165
166  _job_api_resource_data: dict[str, Any]
167  _information_schema_job_metadata: dict[str, Any]
168  project_id: str
169
170  def __init__(
171      self,
172      project_id: str,
173      job_api_resource_data: dict[str, Any],
174      information_schema_job_metadata: dict[str, str],
175  ):
176    super().__init__(project_id)
177    self._job_api_resource_data = job_api_resource_data
178    self._information_schema_job_metadata = (information_schema_job_metadata or
179                                             {})
180
181  @property
182  def full_path(self) -> str:
183    # returns 'https://content-bigquery.googleapis.com/bigquery/v2/
184    # projects/<PROJECT_ID>/jobs/<JOBID>?location=<REGION>'
185    return self._job_api_resource_data.get('selfLink', '')
186
187  @property
188  def id(self) -> str:
189    # returns <PROJECT>:<REGION>.<JobID>
190    return self._job_api_resource_data.get('id', '')
191
192  @property
193  def short_path(self) -> str:
194    # returns <PROJECT>:<REGION>.<JobID>
195    return self.id
196
197  @property
198  def user_email(self) -> str:
199    return self._job_api_resource_data.get('user_email', '')
200
201  @property
202  def _job_configuration(self) -> dict[str, Any]:
203    return self._job_api_resource_data.get('configuration', {})
204
205  @property
206  def _query(self) -> dict[str, Any]:
207    return self._job_configuration.get('query', {})
208
209  @property
210  def _stats(self) -> dict[str, Any]:
211    """Safely access the 'statistics' dictionary."""
212    return self._job_api_resource_data.get('statistics', {})
213
214  @property
215  def _query_stats(self) -> dict[str, Any]:
216    """Safely access the 'statistics.query' dictionary."""
217    return self._stats.get('query', {})
218
219  @property
220  def _query_info(self) -> dict[str, Any]:
221    return self._query_stats.get('queryInfo', {})
222
223  @property
224  def _status(self) -> dict[str, Any]:
225    return self._job_api_resource_data.get('status', {})
226
227  @property
228  def job_type(self) -> str:
229    return self._job_configuration.get('jobType', '')
230
231  @property
232  def query_sql(self) -> str:
233    return self._query.get('query', '')
234
235  @property
236  def use_legacy_sql(self) -> bool:
237    return self._query.get('useLegacySql', False)
238
239  @property
240  def priority(self) -> str:
241    return self._query.get('priority', '')
242
243  @property
244  def edition(self) -> str:
245    edition_value = self._query.get('edition')
246    return str(edition_value) if edition_value else ''
247
248  @property
249  def creation_time(self) -> Optional[int]:
250    time_str = self._stats.get('creationTime')
251    return (int(time_str)
252            if isinstance(time_str, str) and time_str.isdigit() else None)
253
254  @property
255  def start_time(self) -> Optional[int]:
256    time_str = self._stats.get('startTime')
257    return (int(time_str)
258            if isinstance(time_str, str) and time_str.isdigit() else None)
259
260  @property
261  def end_time(self) -> Optional[int]:
262    time_str = self._stats.get('endTime')
263    return (int(time_str)
264            if isinstance(time_str, str) and time_str.isdigit() else None)
265
266  @property
267  def total_bytes_processed(self) -> int:
268    bytes_str = self._stats.get('totalBytesProcessed', '0')
269    return (int(bytes_str)
270            if isinstance(bytes_str, str) and bytes_str.isdigit() else 0)
271
272  @property
273  def total_bytes_billed(self) -> int:
274    bytes_str = self._query_stats.get('totalBytesBilled', '0')
275    return (int(bytes_str)
276            if isinstance(bytes_str, str) and bytes_str.isdigit() else 0)
277
278  @property
279  def total_slot_ms(self) -> int:
280    ms_str = self._stats.get('totalSlotMs', '0')
281    return int(ms_str) if isinstance(ms_str, str) and ms_str.isdigit() else 0
282
283  @property
284  def cache_hit(self) -> bool:
285    return self._query_stats.get('cacheHit') is True
286
287  @property
288  def quota_deferments(self) -> list[str]:
289    deferments_dict = self._stats.get('quotaDeferments', {})
290    if isinstance(deferments_dict, dict):
291      deferment_list = deferments_dict.get('', [])
292      if isinstance(deferment_list, list) and all(
293          isinstance(s, str) for s in deferment_list):
294        return deferment_list
295    return []
296
297  @property
298  def query_plan(self) -> list[dict[str, Any]]:
299    plan = self._query_stats.get('queryPlan', [])
300    return plan if isinstance(plan, list) else []
301
302  @property
303  def total_partitions_processed(self) -> int:
304    partitions_str = self._query_stats.get('totalPartitionsProcessed', '0')
305    return (int(partitions_str) if isinstance(partitions_str, str) and
306            partitions_str.isdigit() else 0)
307
308  @property
309  def referenced_tables(self) -> list[BigQueryTable]:
310    tables_list = self._query_stats.get('referencedTables', [])
311    referenced_tables = []
312    if isinstance(tables_list, list):
313      for item in tables_list:
314        if isinstance(item, dict):
315          project_id = item.get('projectId')
316          dataset_id = item.get('datasetId')
317          table_id = item.get('tableId')
318          if (isinstance(project_id, str) and project_id and
319              isinstance(dataset_id, str) and dataset_id and
320              isinstance(table_id, str) and table_id):
321            referenced_tables.append(
322                BigQueryTable(project_id, dataset_id, table_id))
323    return referenced_tables
324
325  @property
326  def referenced_routines(self) -> list[BigQueryRoutine]:
327    routines_list = self._query_stats.get('referencedRoutines', [])
328    referenced_routines = []
329    if isinstance(routines_list, list):
330      for item in routines_list:
331        if isinstance(item, dict):
332          project_id = item.get('projectId')
333          dataset_id = item.get('datasetId')
334          routine_id = item.get('routineId')
335          if (isinstance(project_id, str) and project_id and
336              isinstance(dataset_id, str) and dataset_id and
337              isinstance(routine_id, str) and routine_id):
338            referenced_routines.append(
339                BigQueryRoutine(project_id, dataset_id, routine_id))
340    return referenced_routines
341
342  @property
343  def num_affected_dml_rows(self) -> int:
344    rows_str = self._query_stats.get('numDmlAffectedRows', '0')
345    return (int(rows_str)
346            if isinstance(rows_str, str) and rows_str.isdigit() else 0)
347
348  @property
349  def dml_stats(self) -> dict[str, int]:
350    stats = self._query_stats.get('dmlStats')
351    if not isinstance(stats, dict):
352      return {}
353    inserted_str = stats.get('insertedRowCount', '0')
354    deleted_str = stats.get('deletedRowCount', '0')
355    updated_str = stats.get('updatedRowCount', '0')
356    return {
357        'insertedRowCount':
358            (int(inserted_str) if isinstance(inserted_str, str) and
359             inserted_str.isdigit() else 0),
360        'deletedRowCount': (int(deleted_str) if isinstance(deleted_str, str) and
361                            deleted_str.isdigit() else 0),
362        'updatedRowCount': (int(updated_str) if isinstance(updated_str, str) and
363                            updated_str.isdigit() else 0),
364    }
365
366  @property
367  def statement_type(self) -> str:
368    stype = self._query_stats.get('statementType', '')
369    return stype if isinstance(stype, str) else ''
370
371  @property
372  def bi_engine_statistics(self) -> dict[str, Any]:
373    stats = self._query_stats.get('biEngineStatistics')
374    if not isinstance(stats, dict):
375      return {}
376    reasons_list = stats.get('accelerationMode', {}).get('biEngineReasons', [])
377    bi_engine_reasons = []
378    if isinstance(reasons_list, list):
379      for item in reasons_list:
380        if isinstance(item, dict):
381          bi_engine_reasons.append({
382              'code': str(item.get('code', '')),
383              'message': item.get('message', ''),
384          })
385    return {
386        'biEngineMode': str(stats.get('biEngineMode', '')),
387        'accelerationMode': str(stats.get('accelerationMode', '')),
388        'biEngineReasons': bi_engine_reasons,
389    }
390
391  @property
392  def vector_search_statistics(self) -> dict[str, Any]:
393    stats = self._query_stats.get('vectorSearchStatistics')
394    if not isinstance(stats, dict):
395      return {}
396    reasons_list = stats.get('indexUnusedReasons', [])
397    index_unused_reasons = []
398    if isinstance(reasons_list, list):
399      for item in reasons_list:
400        if isinstance(item, dict):
401          base_table_data = item.get('baseTable')
402          base_table_obj = None
403          if isinstance(base_table_data, dict):
404            project_id = base_table_data.get('projectId')
405            dataset_id = base_table_data.get('datasetId')
406            table_id = base_table_data.get('tableId')
407            if (isinstance(project_id, str) and project_id and
408                isinstance(dataset_id, str) and dataset_id and
409                isinstance(table_id, str) and table_id):
410              base_table_obj = BigQueryTable(project_id, dataset_id, table_id)
411          index_unused_reasons.append({
412              'code': str(item.get('code', '')),
413              'message': item.get('message', ''),
414              'indexName': item.get('indexName', ''),
415              'baseTable': base_table_obj,
416          })
417    return {
418        'indexUsageMode': str(stats.get('indexUsageMode', '')),
419        'indexUnusedReasons': index_unused_reasons,
420    }
421
422  @property
423  def performance_insights(self) -> dict[str, Any]:
424    insights = self._query_stats.get('performanceInsights')
425    if not isinstance(insights, dict):
426      return {}
427    standalone_list = insights.get('stagePerformanceStandaloneInsights', [])
428    stage_performance_standalone_insights = []
429    if isinstance(standalone_list, list):
430      for item in standalone_list:
431        if isinstance(item, dict):
432          stage_performance_standalone_insights.append({
433              'stageId': item.get('stageId', ''),
434          })
435    change_list = insights.get('stagePerformanceChangeInsights', [])
436    stage_performance_change_insights = []
437    if isinstance(change_list, list):
438      for item in change_list:
439        if isinstance(item, dict):
440          stage_performance_change_insights.append({
441              'stageId': item.get('stageId', ''),
442          })
443    avg_ms_str = insights.get('avgPreviousExecutionMs', '0')
444    return {
445        'avgPreviousExecutionMs':
446            (int(avg_ms_str)
447             if isinstance(avg_ms_str, str) and avg_ms_str.isdigit() else 0),
448        'stagePerformanceStandaloneInsights':
449            (stage_performance_standalone_insights),
450        'stagePerformanceChangeInsights': stage_performance_change_insights,
451    }
452
453  @property
454  def optimization_details(self) -> Any:
455    return self._query_info.get('optimizationDetails')
456
457  @property
458  def export_data_statistics(self) -> dict[str, int]:
459    stats = self._query_stats.get('exportDataStatistics')
460    if not isinstance(stats, dict):
461      return {}
462    file_count_str = stats.get('fileCount', '0')
463    row_count_str = stats.get('rowCount', '0')
464    return {
465        'fileCount': (int(file_count_str) if isinstance(file_count_str, str) and
466                      file_count_str.isdigit() else 0),
467        'rowCount': (int(row_count_str) if isinstance(row_count_str, str) and
468                     row_count_str.isdigit() else 0),
469    }
470
471  @property
472  def load_query_statistics(self) -> dict[str, int]:
473    stats = self._query_stats.get('loadQueryStatistics')
474    if not isinstance(stats, dict):
475      return {}
476    input_files_str = stats.get('inputFiles', '0')
477    input_bytes_str = stats.get('inputFileBytes', '0')
478    output_rows_str = stats.get('outputRows', '0')
479    output_bytes_str = stats.get('outputBytes', '0')
480    bad_records_str = stats.get('badRecords', '0')
481    return {
482        'inputFiles':
483            (int(input_files_str) if isinstance(input_files_str, str) and
484             input_files_str.isdigit() else 0),
485        'inputFileBytes':
486            (int(input_bytes_str) if isinstance(input_bytes_str, str) and
487             input_bytes_str.isdigit() else 0),
488        'outputRows':
489            (int(output_rows_str) if isinstance(output_rows_str, str) and
490             output_rows_str.isdigit() else 0),
491        'outputBytes':
492            (int(output_bytes_str) if isinstance(output_bytes_str, str) and
493             output_bytes_str.isdigit() else 0),
494        'badRecords':
495            (int(bad_records_str) if isinstance(bad_records_str, str) and
496             bad_records_str.isdigit() else 0),
497    }
498
499  @property
500  def spark_statistics(self) -> dict[str, Any]:
501    stats = self._query_stats.get('sparkStatistics')
502    if not isinstance(stats, dict):
503      return {}
504    logging_info_dict = stats.get('loggingInfo', {})
505    logging_info = ({
506        'resourceType': logging_info_dict.get('resourceType', ''),
507        'projectId': logging_info_dict.get('projectId', ''),
508    } if isinstance(logging_info_dict, dict) else {})
509    return {
510        'endpoints': stats.get('endpoints', {}),
511        'sparkJobId': stats.get('sparkJobId', ''),
512        'sparkJobLocation': stats.get('sparkJobLocation', ''),
513        'kmsKeyName': stats.get('kmsKeyName', ''),
514        'gcsStagingBucket': stats.get('gcsStagingBucket', ''),
515        'loggingInfo': logging_info,
516    }
517
518  @property
519  def transferred_bytes(self) -> int:
520    bytes_str = self._query_stats.get('transferredBytes', '0')
521    return (int(bytes_str)
522            if isinstance(bytes_str, str) and bytes_str.isdigit() else 0)
523
524  @property
525  def reservation_id(self) -> str:
526    res_id = self._stats.get('reservation_id', '')
527    return res_id if isinstance(res_id, str) else ''
528
529  @property
530  def reservation_admin_project_id(self) -> Optional[str]:
531    if not self.reservation_id:
532      return None
533    try:
534      parts = self.reservation_id.split('/')
535      if parts[0] == 'projects' and len(parts) >= 2:
536        return parts[1]
537      else:
538        logging.warning(
539            'Could not parse project ID from reservation_id: %s',
540            self.reservation_id,
541        )
542        return None
543    except (IndexError, AttributeError):
544      logging.warning(
545          'Could not parse project ID from reservation_id: %s',
546          self.reservation_id,
547      )
548      return None
549
550  @property
551  def num_child_jobs(self) -> int:
552    num_str = self._stats.get('numChildJobs', '0')
553    return int(num_str) if isinstance(num_str, str) and num_str.isdigit() else 0
554
555  @property
556  def parent_job_id(self) -> str:
557    parent_id = self._stats.get('parentJobId', '')
558    return parent_id if isinstance(parent_id, str) else ''
559
560  @property
561  def row_level_security_applied(self) -> bool:
562    rls_stats = self._stats.get('RowLevelSecurityStatistics', {})
563    return (rls_stats.get('rowLevelSecurityApplied') is True if isinstance(
564        rls_stats, dict) else False)
565
566  @property
567  def data_masking_applied(self) -> bool:
568    masking_stats = self._stats.get('dataMaskingStatistics', {})
569    return (masking_stats.get('dataMaskingApplied') is True if isinstance(
570        masking_stats, dict) else False)
571
572  @property
573  def session_id(self) -> str:
574    session_info = self._stats.get('sessionInfo', {})
575    session_id_val = (session_info.get('sessionId', '') if isinstance(
576        session_info, dict) else '')
577    return session_id_val if isinstance(session_id_val, str) else ''
578
579  @property
580  def final_execution_duration_ms(self) -> int:
581    duration_str = self._stats.get('finalExecutionDurationMs', '0')
582    return (int(duration_str)
583            if isinstance(duration_str, str) and duration_str.isdigit() else 0)
584
585  @property
586  def job_state(self) -> str:
587    state = self._status.get('state', '')
588    return state if isinstance(state, str) else ''
589
590  @property
591  def job_error_result(self) -> dict[str, Optional[str]]:
592    error_result = self._status.get('errorResult')
593    if not isinstance(error_result, dict):
594      return {}
595    return {
596        'reason': error_result.get('reason'),
597        'location': error_result.get('location'),
598        'debugInfo': error_result.get('debugInfo'),
599        'message': error_result.get('message'),
600    }
601
602  @property
603  def job_errors(self) -> list[dict[str, Optional[str]]]:
604    errors_list = self._status.get('errors', [])
605    errors_iterable = []
606    if isinstance(errors_list, list):
607      for item in errors_list:
608        if isinstance(item, dict):
609          errors_iterable.append({
610              'reason': item.get('reason'),
611              'location': item.get('location'),
612              'debugInfo': item.get('debugInfo'),
613              'message': item.get('message'),
614          })
615    return errors_iterable
616
617  @property
618  def materialized_view_statistics(self) -> dict[str, Any]:
619    stats_list = self._query_stats.get('materializedViewStatistics')
620    materialized_view = []
621    if isinstance(stats_list, list):
622      for item in stats_list:
623        if isinstance(item, dict):
624          table_ref_data = item.get('tableReference')
625          table_ref_obj = None
626          if isinstance(table_ref_data, dict):
627            project_id = table_ref_data.get('projectId')
628            dataset_id = table_ref_data.get('datasetId')
629            table_id = table_ref_data.get('tableId')
630            if (isinstance(project_id, str) and project_id and
631                isinstance(dataset_id, str) and dataset_id and
632                isinstance(table_id, str) and table_id):
633              table_ref_obj = BigQueryTable(project_id, dataset_id, table_id)
634          chosen = item.get('chosen') is True
635          saved_str = item.get('estimatedBytesSaved', '0')
636          estimated_bytes_saved = (int(saved_str)
637                                   if isinstance(saved_str, str) and
638                                   saved_str.isdigit() else 0)
639          rejected_reason = str(item.get('rejectedReason', ''))
640          materialized_view.append({
641              'chosen': chosen,
642              'estimatedBytesSaved': estimated_bytes_saved,
643              'rejectedReason': rejected_reason,
644              'tableReference': table_ref_obj,
645          })
646    return {'materializedView': materialized_view}
647
648  @property
649  def metadata_cache_statistics(self) -> dict[str, Any]:
650    stats_list = self._query_stats.get('metadataCacheStatistics')
651    metadata_cache = []
652    if isinstance(stats_list, list):
653      for item in stats_list:
654        if isinstance(item, dict):
655          table_ref_data = item.get('tableReference')
656          table_ref_obj = None
657          if isinstance(table_ref_data, dict):
658            project_id = table_ref_data.get('projectId')
659            dataset_id = table_ref_data.get('datasetId')
660            table_id = table_ref_data.get('tableId')
661            if (isinstance(project_id, str) and project_id and
662                isinstance(dataset_id, str) and dataset_id and
663                isinstance(table_id, str) and table_id):
664              table_ref_obj = BigQueryTable(project_id, dataset_id, table_id)
665          metadata_cache.append({
666              'explanation': item.get('explanation', ''),
667              'unusedReason': str(item.get('unusedReason', '')),
668              'tableReference': table_ref_obj,
669          })
670    return {'tableMetadataCacheUsage': metadata_cache}
671
672  # Properties derived from _information_schema_job_metadata
673  @property
674  def information_schema_user_email(self) -> str | None:
675    if not self._information_schema_job_metadata:
676      return C_NOT_AVAILABLE
677    return self._information_schema_job_metadata.get('user_email')
678
679  @property
680  def information_schema_start_time_str(self) -> str | None:
681    if not self._information_schema_job_metadata:
682      return C_NOT_AVAILABLE
683    return self._information_schema_job_metadata.get('start_time_str')
684
685  @property
686  def information_schema_end_time_str(self) -> str | None:
687    if not self._information_schema_job_metadata:
688      return C_NOT_AVAILABLE
689    return self._information_schema_job_metadata.get('end_time_str')
690
691  @property
692  def information_schema_query(self) -> str | None:
693    if not self._information_schema_job_metadata:
694      return C_NOT_AVAILABLE
695    return self._information_schema_job_metadata.get('query')
696
697  @property
698  def information_schema_total_modified_partitions(self) -> Union[int, str]:
699    """The total number of partitions the job modified.
700
701    This field is populated for LOAD and QUERY jobs.
702    """
703    if not self._information_schema_job_metadata:
704      return C_NOT_AVAILABLE
705    try:
706      total_modified_partitions = self._information_schema_job_metadata[
707          'total_modified_partitions']
708      return total_modified_partitions
709    except KeyError:
710      return C_NOT_AVAILABLE
711
712  @property
713  def information_schema_resource_warning(self) -> str:
714    """The warning message that appears if the resource usage during query
715
716    processing is above the internal threshold of the system.
717    """
718    if not self._information_schema_job_metadata:
719      return C_NOT_AVAILABLE
720    try:
721      resource_warning = self._information_schema_job_metadata['query_info'][
722          'resource_warning']
723      return resource_warning
724    except KeyError:
725      return C_NOT_AVAILABLE
726
727  @property
728  def information_schema_normalized_literals(self) -> str:
729    """Contains the hashes of the query."""
730    try:
731      query_hashes = self._information_schema_job_metadata['query_info'][
732          'query_hashes']['normalized_literals']
733      return query_hashes
734    except KeyError:
735      return C_NOT_AVAILABLE