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, organization_id, raise_error_if_fails=False)
 89    return policy
 90  except utils.GcpApiError as err:
 91    if (
 92      "doesn't have access to" in err.message.lower() or 'denied on resource' in err.message.lower()
 93    ):
 94      op.info(
 95        'User does not have access to the organization policy. Investigation'
 96        ' completeness and accuracy might depend on the presence of'
 97        ' organization level permissions.'
 98      )
 99    return None
100  finally:
101    root_logger.setLevel(original_level)

104def check_permissions_for_principal(
105  policy: PolicyObject, principal: str, permissions_to_check: Set[str]
106) -> 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) for permission in permissions_to_check
114  }

117def get_missing_permissions(
118  required_permissions: Set[str], actual_permissions: Dict[str, bool]
119) -> 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 {perm for perm in required_permissions if not actual_permissions.get(perm)}

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

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

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