#!/usr/bin/env python
"""Python script for interacting with BigQuery."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import datetime
import functools
import json
import logging
import os
import re
import sys
from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Type

from absl import app
from absl import flags
import yaml

import table_formatter
import bq_utils
from clients import utils as bq_client_utils
from frontend import utils_flags
from frontend import utils_formatting
from utils import bq_consts
from utils import bq_error
from utils import bq_id_utils
from pyglib import stringutil

# pylint: disable=g-multiple-import
if sys.version_info < (3, 11):
  from typing_extensions import TypedDict, NotRequired  # pylint: disable=g-import-not-at-top
else:
  from typing import TypedDict, NotRequired  # pylint: disable=g-import-not-at-top
# pylint: enable=g-multiple-import

FLAGS = flags.FLAGS

PARQUET_LIST_INFERENCE_DESCRIPTION = (
    'Use schema inference specifically for Parquet LIST logical type.\n It'
    ' checks whether the LIST node is in the standard form as documented in:\n'
    ' https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#lists\n'
    '  <optional | required> group <name> (LIST) {\n    repeated group list {\n'
    '      <optional | required> <element-type> element;\n    }\n  }\n Returns'
    ' the "element" node in list_element_node. The corresponding field for the'
    ' LIST node in the converted schema is treated as if the node has the'
    ' following schema:\n repeated <element-type> <name>\n This means nodes'
    ' "list" and "element" are omitted.\n\n Otherwise, the LIST node must be in'
    ' one of the forms described by the backward-compatibility rules as'
    ' documented in:\n'
    ' https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#backward-compatibility-rules\n'
    ' <optional | required> group <name> (LIST) {\n   repeated <element-type>'
    ' <element-name>\n }\n Returns the <element-name> node in'
    ' list_element_node. The corresponding field for the LIST node in the'
    ' converted schema is treated as if the node has the following schema:\n'
    ' repeated <element-type> <name>\n This means the element node is omitted.'
)

CONNECTION_ID_PATTERN = re.compile(r'[\w-]+')
_RANGE_PATTERN = re.compile(r'^\[(\S+.+\S+), (\S+.+\S+)\)$')

_PARAMETERS_KEY = 'parameters'
_DEFAULT_STORAGE_LOCATION_URI_KEY = 'defaultStorageLocationUri'
_STORAGE_DESCRIPTOR_KEY = 'storageDescriptor'
_CONNECTION_ID_KEY = 'connectionId'

_DELIMITER_MAP = {
    'tab': '\t',
    '\\t': '\t',
}
_DDL_OPERATION_MAP = {
    'SKIP': 'Skipped',
    'CREATE': 'Created',
    'REPLACE': 'Replaced',
    'ALTER': 'Altered',
    'DROP': 'Dropped',
}


def ValidateGlobalFlags():
  """Validate combinations of global flag values."""
  if FLAGS.service_account and FLAGS.use_gce_service_account:
    raise app.UsageError(
        'Cannot specify both --service_account and --use_gce_service_account.'
    )


def ValidateAtMostOneSelected(*args: Any) -> bool:
  """Validates that at most one of the argument flags is selected.

  Args:
    *args: Each flag to be tested parsed in as a separate arg.

  Returns:
    True if more than 1 flag was selected, False if 1 or 0 were selected.
  """
  count = 0
  for arg in args:
    if arg:
      count += 1
  return count > 1


def ValidateAtMostOneSelectedAllowsDefault(*args: Any) -> bool:
  """Validates that at most one of the argument flags is selected.

    if the arg exists but the value is the default value,
    then it won't be counted. This is uself when users want to clear the
    value while setting another value. For example, 'update --arg1=0 --arg2=100'
    when arg1 and arg2 shouldn't coexist.

  Args:
    *args: Each flag to be tested parsed in as a separate arg.

  Returns:
    True if more than 1 flag was selected, False if 1 or 0 were selected.
  """
  count = 0
  for arg in args:
    if arg and type(arg)() != arg:
      count += 1
  return count > 1


def ProcessSource(description: str, source: str) -> Tuple[Any, Any]:
  """Process "source" parameter used for bq update and bq mk command.

  Args:
    description: Description of the dataset.
    source: source file path attached by "--source" parameter.

  Returns:
    new description if the source file updates the description, otherwise return
    the original description.
    acl if the source file updates the acl, otherwise return None.
  """
  acl = None
  if source is None:
    return (description, acl)
  if not os.path.exists(source):
    raise app.UsageError('Source file not found: %s' % (source,))
  if not os.path.isfile(source):
    raise app.UsageError('Source path is not a file: %s' % (source,))
  with open(source) as f:
    try:
      payload = json.load(f)
      if 'description' in payload:
        description = payload['description']
        logging.debug(
            'Both source file and description flag exist, using the value in'
            ' the source file.'
        )
      if 'access' in payload:
        acl = payload['access']
    except ValueError as e:
      raise app.UsageError(
          'Error decoding JSON schema from file %s: %s' % (source, e)
      )

  return (description, acl)


def PrintDryRunInfo(job):
  """Prints the dry run info."""
  if FLAGS.format in ['prettyjson', 'json']:
    bq_utils.PrintFormattedJsonObject(job)
    return
  # TODO: b/422281423 - Revert this check once server returns this field for
  # all dry run queries again.
  if 'totalBytesProcessed' not in job['statistics']['query']:
    print(
        'Query successfully validated. No information about number of bytes'
        ' processed. To see the full details of the job, run this query with'
        ' `--format=json` or `--format=prettyjson`.'
    )
    return
  num_bytes = job['statistics']['query']['totalBytesProcessed']
  num_bytes_accuracy = job['statistics']['query'].get(
      'totalBytesProcessedAccuracy', 'PRECISE'
  )
  if FLAGS.format == 'csv':
    print(num_bytes)
  else:
    if job['statistics']['query'].get('statementType', '') == 'LOAD_DATA':
      print(
          'Query successfully validated. Assuming the files are not modified, '
          'running this query will process %s files loading %s bytes of data.'
          % (
              job['statistics']['query']['loadQueryStatistics']['inputFiles'],
              job['statistics']['query']['loadQueryStatistics'][
                  'inputFileBytes'
              ],
          )
      )
    elif num_bytes_accuracy == 'PRECISE':
      print(
          'Query successfully validated. Assuming the tables are not modified, '
          'running this query will process %s bytes of data.' % (num_bytes,)
      )
    elif num_bytes_accuracy == 'LOWER_BOUND':
      print(
          'Query successfully validated. Assuming the tables are not modified, '
          'running this query will process lower bound of %s bytes of data.'
          % (num_bytes,)
      )
    elif num_bytes_accuracy == 'UPPER_BOUND':
      print(
          'Query successfully validated. Assuming the tables are not modified, '
          'running this query will process upper bound of %s bytes of data.'
          % (num_bytes,)
      )
    else:
      if job['statistics']['query']['statementType'] == 'CREATE_MODEL':
        print(
            'Query successfully validated. The number of bytes that will '
            'be processed by this query cannot be calculated automatically. '
            'More information about this can be seen in '
            'https://cloud.google.com/bigquery-ml/pricing#dry_run'
        )
      else:
        print(
            'Query successfully validated. Assuming the tables are not '
            'modified, running this query will process %s of data and the '
            'accuracy is unknown because of federated tables or clustered '
            'tables.' % (num_bytes,)
        )


def RawInput(message: str) -> str:
  try:
    return input(message)
  except EOFError:
    if sys.stdin.isatty():
      print('\nGot EOF; exiting.')
    else:
      print('\nGot EOF; exiting. Is your input from a terminal?')
    sys.exit(1)


def PromptWithDefault(message: str) -> str:
  """Prompts user with message, return key pressed or '' on enter."""
  if FLAGS.headless:
    print('Running --headless, accepting default for prompt: %s' % (message,))
    return ''
  return RawInput(message).lower()


def PromptYN(message: str) -> Optional[str]:
  """Prompts user with message, returning the key 'y', 'n', or '' on enter."""
  response = None
  while response not in ['y', 'n', '']:
    response = PromptWithDefault(message)
  return response


def NormalizeFieldDelimiter(field_delimiter: str) -> str:
  """Validates and returns the correct field_delimiter."""
  # The only non-string delimiter we allow is None, which represents
  # no field delimiter specified by the user.
  if field_delimiter is None:
    return field_delimiter

  # Allow TAB and \\t substitution.
  key = field_delimiter.lower()
  return _DELIMITER_MAP.get(key, field_delimiter)


def ValidateHivePartitioningOptions(hive_partitioning_mode):
  """Validates the string provided is one the API accepts.

  Should not receive None as an input, since that will fail the comparison.
  Args:
    hive_partitioning_mode: String representing which hive partitioning mode is
      requested.  Only 'AUTO' and 'STRINGS' are supported.
  """
  if hive_partitioning_mode not in ['AUTO', 'STRINGS', 'CUSTOM']:
    raise app.UsageError(
        'Only the following hive partitioning modes are supported: "AUTO", '
        '"STRINGS" and "CUSTOM"'
    )


def ParseLabels(labels: List[str]) -> Dict[str, str]:
  """Parses a list of user-supplied strings representing labels.

  Args:
    labels: A list of user-supplied strings representing labels.  It is expected
      to be in the format "key:value".

  Returns:
    A dict mapping label keys to label values.

  Raises:
    UsageError: Incorrect label arguments were supplied.
  """
  labels_dict = {}
  for key_value in labels:
    k, _, v = key_value.partition(':')
    k = k.strip()
    if k in labels_dict:
      raise app.UsageError('Cannot specify label key "%s" multiple times' % k)
    if k.strip():
      labels_dict[k.strip()] = v.strip()
  return labels_dict


def IsRangeBoundaryUnbounded(value: str) -> bool:
  return value.upper() == 'UNBOUNDED' or value.upper() == 'NULL'


def ParseRangeString(value: str) -> Optional[Tuple[str, str]]:
  match = _RANGE_PATTERN.match(value)
  if not match:
    return None
  start, end = match.groups()
  return start, end


class TablePrinter(object):
  """Base class for printing a table, with a default implementation."""

  def __init__(self, **kwds):
    super(TablePrinter, self).__init__()
    # Most extended classes will require state.
    for key, value in kwds.items():
      setattr(self, key, value)

  @staticmethod
  def _ValidateFields(fields, formatter):
    if isinstance(formatter, table_formatter.CsvFormatter):
      for field in fields:
        if field['type'].upper() == 'RECORD':
          raise app.UsageError(
              (
                  'Error printing table: Cannot print record '
                  'field "%s" in CSV format.'
              )
              % field['name']
          )
        if field.get('mode', 'NULLABLE').upper() == 'REPEATED':
          raise app.UsageError(
              (
                  'Error printing table: Cannot print repeated '
                  'field "%s" in CSV format.'
              )
              % (field['name'])
          )

  @staticmethod
  def _NormalizeRecord(field, value, use_full_timestamp):
    """Returns bq-specific formatting of a RECORD type."""
    result = collections.OrderedDict()
    for subfield, subvalue in zip(field.get('fields', []), value):
      result[subfield.get('name', '')] = TablePrinter.NormalizeField(
          subfield, subvalue, use_full_timestamp
      )
    return result

  @staticmethod
  def _NormalizeTimestamp(unused_field, value, use_full_timestamp):
    """Returns bq-specific formatting of a TIMESTAMP type."""
    try:
      if use_full_timestamp:
        return value
      else:
        date = datetime.datetime.fromtimestamp(
            0, tz=datetime.timezone.utc
        ) + datetime.timedelta(seconds=float(value))
        # Remove the extra timezone info "+00:00" at the end of the date.
        date = date.replace(tzinfo=None)
        # Our goal is the equivalent of '%Y-%m-%d %H:%M:%S' via strftime but that
        # doesn't work for dates with years prior to 1900.  Instead we zero out
        # fractional seconds then call isoformat with a space separator.
        date = date.replace(microsecond=0)
        return date.isoformat(' ')
    except (ValueError, OverflowError):
      return '<date out of range for display>'

  @staticmethod
  def _NormalizeRange(field, value, use_full_timestamp):
    """Returns bq-specific formatting of a RANGE type."""
    parsed = ParseRangeString(value)
    if parsed is None:
      return '<invalid range>'
    start, end = parsed

    if field.get('rangeElementType').get('type').upper() != 'TIMESTAMP':
      start = start.upper() if IsRangeBoundaryUnbounded(start) else start
      end = end.upper() if IsRangeBoundaryUnbounded(end) else end
      return '[%s, %s)' % (start, end)

    if IsRangeBoundaryUnbounded(start):
      normalized_start = start.upper()
    else:
      normalized_start = TablePrinter._NormalizeTimestamp(
          field, start, use_full_timestamp
      )
    if IsRangeBoundaryUnbounded(end):
      normalized_end = end.upper()
    else:
      normalized_end = TablePrinter._NormalizeTimestamp(
          field, end, use_full_timestamp
      )
    return '[%s, %s)' % (normalized_start, normalized_end)

  @staticmethod
  def NormalizeField(field, value, use_full_timestamp: bool):
    """Returns bq-specific formatting of a field."""
    if value is None:
      return None
    if field.get('mode', '').upper() == 'REPEATED':
      return [
          TablePrinter._NormalizeSingleValue(field, value, use_full_timestamp)
          for value in value
      ]
    return TablePrinter._NormalizeSingleValue(field, value, use_full_timestamp)

  @staticmethod
  def _NormalizeSingleValue(field, value, use_full_timestamp: bool):
    """Returns formatting of a single field value."""
    if field.get('type', '').upper() == 'RECORD':
      return TablePrinter._NormalizeRecord(field, value, use_full_timestamp)
    elif field.get('type', '').upper() == 'TIMESTAMP':
      return TablePrinter._NormalizeTimestamp(field, value, use_full_timestamp)
    elif field.get('type', '').upper() == 'RANGE':
      return TablePrinter._NormalizeRange(field, value, use_full_timestamp)
    return value

  @staticmethod
  def MaybeConvertToJson(value):
    """Converts dicts and lists to JSON; returns everything else as-is."""
    if isinstance(value, dict) or isinstance(value, list):
      return json.dumps(value, separators=(',', ':'), ensure_ascii=False)
    return value

  @staticmethod
  def FormatRow(fields, row, formatter, use_full_timestamp: bool):
    """Convert fields in a single row to bq-specific formatting."""
    values = [
        TablePrinter.NormalizeField(field, value, use_full_timestamp)
        for field, value in zip(fields, row)
    ]
    # Convert complex values to JSON if we're not already outputting as such.
    if not isinstance(formatter, table_formatter.JsonFormatter):
      values = map(TablePrinter.MaybeConvertToJson, values)
    # Convert NULL values to strings for CSV and non-JSON formats.
    if isinstance(formatter, table_formatter.CsvFormatter):
      values = ['' if value is None else value for value in values]
    elif not isinstance(formatter, table_formatter.JsonFormatter):
      values = ['NULL' if value is None else value for value in values]
    return values

  def PrintTable(self, fields, rows, use_full_timestamp: bool):
    formatter = utils_flags.get_formatter_from_flags(secondary_format='pretty')
    self._ValidateFields(fields, formatter)
    formatter.AddFields(fields)
    formatter.AddRows(
        TablePrinter.FormatRow(fields, row, formatter, use_full_timestamp)
        for row in rows
    )
    formatter.Print()


def CreateExternalTableDefinition(
    source_format,
    source_uris,
    schema,
    autodetect,
    connection_id=None,
    ignore_unknown_values=False,
    hive_partitioning_mode=None,
    hive_partitioning_source_uri_prefix=None,
    require_hive_partition_filter=None,
    use_avro_logical_types=False,
    parquet_enum_as_string=False,
    parquet_enable_list_inference=False,
    metadata_cache_mode=None,
    object_metadata=None,
    preserve_ascii_control_characters=False,
    reference_file_schema_uri=None,
    encoding=None,
    file_set_spec_type=None,
    null_marker=None,
    null_markers=None,
    time_zone=None,
    date_format=None,
    datetime_format=None,
    time_format=None,
    timestamp_format=None,
    source_column_match=None,
    parquet_map_target_type=None,
    timestamp_target_precision=None,
):
  """Creates an external table definition with the given URIs and the schema.

  Arguments:
    source_format: Format of source data. For CSV files, specify 'CSV'. For
      Google spreadsheet files, specify 'GOOGLE_SHEETS'. For newline-delimited
      JSON, specify 'NEWLINE_DELIMITED_JSON'. For Cloud Datastore backup,
      specify 'DATASTORE_BACKUP'. For Avro files, specify 'AVRO'. For Orc files,
      specify 'ORC'. For Parquet files, specify 'PARQUET'. For Iceberg tables,
      specify 'ICEBERG'.
    source_uris: Comma separated list of URIs that contain data for this table.
    schema: Either an inline schema or path to a schema file.
    autodetect: Indicates if format options, compression mode and schema be auto
      detected from the source data. True - means that autodetect is on, False
      means that it is off. None means format specific default: - For CSV it
      means autodetect is OFF - For JSON it means that autodetect is ON. For
      JSON, defaulting to autodetection is safer because the only option
      autodetected is compression. If a schema is passed, then the user-supplied
      schema is used.
    connection_id: The user flag with the same name defined for the _Load
      BigqueryCmd
    ignore_unknown_values:  Indicates if BigQuery should allow extra values that
      are not represented in the table schema. If true, the extra values are
      ignored. If false, records with extra columns are treated as bad records,
      and if there are too many bad records, an invalid error is returned in the
      job result. The default value is false. The sourceFormat property
      determines what BigQuery treats as an extra value: - CSV: Trailing columns
      - JSON: Named values that don't match any column names.
    hive_partitioning_mode: Enables hive partitioning.  AUTO indicates to
      perform automatic type inference.  STRINGS indicates to treat all hive
      partition keys as STRING typed.  No other values are accepted.
    hive_partitioning_source_uri_prefix: Shared prefix for all files until hive
      partitioning encoding begins.
    require_hive_partition_filter: The user flag with the same name defined for
      the _Load BigqueryCmd
    use_avro_logical_types: The user flag with the same name defined for the
      _Load BigqueryCmd
    parquet_enum_as_string: The user flag with the same name defined for the
      _Load BigqueryCmd
    parquet_enable_list_inference: The user flag with the same name defined for
      the _Load BigqueryCmd
    metadata_cache_mode: Enables metadata cache for an external table with a
      connection. Specify 'AUTOMATIC' to automatically refresh the cached
      metadata. Specify 'MANUAL' to stop the automatic refresh.
    object_metadata: Object Metadata Type.
    preserve_ascii_control_characters: The user flag with the same name defined
      for the _Load BigqueryCmd
    reference_file_schema_uri: The user flag with the same name defined for the
      _Load BigqueryCmd
    encoding: Encoding types for CSV files. Available options are: 'UTF-8',
      'ISO-8859-1', 'UTF-16BE', 'UTF-16LE', 'UTF-32BE', and 'UTF-32LE'. The
      default value is 'UTF-8'.
    file_set_spec_type: Set how to discover files given source URIs. Specify
      'FILE_SYSTEM_MATCH' (default behavior) to expand source URIs by listing
      files from the underlying object store. Specify
      'NEW_LINE_DELIMITED_MANIFEST' to parse the URIs as new line delimited
      manifest files, where each line contains a URI (No wild-card URIs are
      supported).
    null_marker: Specifies a string that represents a null value in a CSV file.
    null_markers: Specifies a list of strings that represent null values in a
      CSV file.
    time_zone: Specifies the time zone for a CSV or JSON file.
    date_format: Specifies the date format for a CSV or JSON file.
    datetime_format: Specifies the datetime format for a CSV or JSON file.
    time_format: Specifies the time format for a CSV or JSON file.
    timestamp_format: Specifies the timestamp format for a CSV or JSON file.
    source_column_match: Controls the strategy used to match loaded columns to
      the schema.
    parquet_map_target_type: Indicate the target type for parquet maps. If
      unspecified, we represent parquet maps as map {repeated key_value {key,
      value}}. This option can simplify this by omiting the key_value record if
      it's equal to ARRAY_OF_STRUCT.
    timestamp_target_precision: Precision (maximum number of total digits in
      base 10) for seconds of TIMESTAMP type.

  Returns:
    A python dictionary that contains a external table definition for the given
    format with the most common options set.
  """
  try:
    supported_formats = [
        'CSV',
        'NEWLINE_DELIMITED_JSON',
        'DATASTORE_BACKUP',
        'DELTA_LAKE',
        'AVRO',
        'ORC',
        'PARQUET',
        'GOOGLE_SHEETS',
        'ICEBERG',
    ]

    if source_format not in supported_formats:
      raise app.UsageError('%s is not a supported format.' % source_format)

    external_table_def = {'sourceFormat': source_format}
    if file_set_spec_type is not None:
      external_table_def['fileSetSpecType'] = file_set_spec_type
    if metadata_cache_mode is not None:
      external_table_def['metadataCacheMode'] = metadata_cache_mode
    if time_zone is not None:
      external_table_def['timeZone'] = time_zone
    if date_format is not None:
      external_table_def['dateFormat'] = date_format
    if datetime_format is not None:
      external_table_def['datetimeFormat'] = datetime_format
    if time_format is not None:
      external_table_def['timeFormat'] = time_format
    if timestamp_format is not None:
      external_table_def['timestampFormat'] = timestamp_format
    if object_metadata is not None:
      supported_obj_metadata_types = ['DIRECTORY', 'SIMPLE']

      if object_metadata not in supported_obj_metadata_types:
        raise app.UsageError(
            '%s is not a supported Object Metadata Type.' % object_metadata
        )

      external_table_def['sourceFormat'] = None
      external_table_def['objectMetadata'] = object_metadata
    if timestamp_target_precision is not None:
      external_table_def['timestampTargetPrecision'] = timestamp_target_precision

    if external_table_def['sourceFormat'] == 'CSV':
      if autodetect:
        external_table_def['autodetect'] = True
        external_table_def['csvOptions'] = yaml.safe_load("""
            {
                "quote": '"',
                "encoding": "UTF-8"
            }
        """)
      else:
        external_table_def['csvOptions'] = yaml.safe_load("""
            {
                "allowJaggedRows": false,
                "fieldDelimiter": ",",
                "allowQuotedNewlines": false,
                "quote": '"',
                "skipLeadingRows": 0,
                "encoding": "UTF-8"
            }
        """)
      external_table_def['csvOptions'][
          'preserveAsciiControlCharacters'
      ] = preserve_ascii_control_characters
      external_table_def['csvOptions']['encoding'] = encoding or 'UTF-8'
      if null_marker is not None:
        external_table_def['csvOptions']['nullMarker'] = null_marker
      if null_markers is not None:
        external_table_def['csvOptions']['nullMarkers'] = null_markers
      if source_column_match is not None:
        external_table_def['csvOptions'][
            'sourceColumnMatch'
        ] = source_column_match
    elif external_table_def['sourceFormat'] == 'NEWLINE_DELIMITED_JSON':
      if autodetect is None or autodetect:
        external_table_def['autodetect'] = True
      external_table_def['jsonOptions'] = {'encoding': encoding or 'UTF-8'}
    elif external_table_def['sourceFormat'] == 'GOOGLE_SHEETS':
      if autodetect is None or autodetect:
        external_table_def['autodetect'] = True
      else:
        external_table_def['googleSheetsOptions'] = yaml.safe_load("""
            {
                "skipLeadingRows": 0
            }
        """)
    elif external_table_def['sourceFormat'] == 'AVRO':
      external_table_def['avroOptions'] = {
          'useAvroLogicalTypes': use_avro_logical_types
      }
      if reference_file_schema_uri is not None:
        external_table_def['referenceFileSchemaUri'] = reference_file_schema_uri
    elif external_table_def['sourceFormat'] == 'PARQUET':
      external_table_def['parquetOptions'] = {
          'enumAsString': parquet_enum_as_string,
          'enableListInference': parquet_enable_list_inference,
          'mapTargetType': parquet_map_target_type,
      }
      if reference_file_schema_uri is not None:
        external_table_def['referenceFileSchemaUri'] = reference_file_schema_uri
    elif external_table_def['sourceFormat'] == 'ORC':
      if reference_file_schema_uri is not None:
        external_table_def['referenceFileSchemaUri'] = reference_file_schema_uri
    elif (
        external_table_def['sourceFormat'] == 'ICEBERG'
        or external_table_def['sourceFormat'] == 'DELTA_LAKE'
    ):
      source_format = (
          'Iceberg'
          if external_table_def['sourceFormat'] == 'ICEBERG'
          else 'Delta Lake'
      )
      if autodetect is not None and not autodetect or schema:
        raise app.UsageError(
            'Cannot create %s table from user-specified schema.'
            % (source_format,)
        )
      # Always autodetect schema for ICEBERG from manifest
      external_table_def['autodetect'] = True
      if len(source_uris.split(',')) != 1:
        raise app.UsageError(
            'Must provide only one source_uri for %s table.' % (source_format,)
        )


    if ignore_unknown_values:
      external_table_def['ignoreUnknownValues'] = True


    if hive_partitioning_mode is not None:
      ValidateHivePartitioningOptions(hive_partitioning_mode)
      hive_partitioning_options = {}
      hive_partitioning_options['mode'] = hive_partitioning_mode
      if hive_partitioning_source_uri_prefix is not None:
        hive_partitioning_options['sourceUriPrefix'] = (
            hive_partitioning_source_uri_prefix
        )
      external_table_def['hivePartitioningOptions'] = hive_partitioning_options
      if require_hive_partition_filter:
        hive_partitioning_options['requirePartitionFilter'] = True

    if schema:
      fields = bq_client_utils.ReadSchema(schema)
      external_table_def['schema'] = {'fields': fields}

    if connection_id:
      external_table_def['connectionId'] = connection_id

    external_table_def['sourceUris'] = source_uris.split(',')

    return external_table_def

  except ValueError as e:
    raise app.UsageError(
        'Error occurred while creating table definition: %s' % e
    )


def GetExternalDataConfig(
    file_path_or_simple_spec,
    use_avro_logical_types=False,
    parquet_enum_as_string=False,
    parquet_enable_list_inference=False,
    metadata_cache_mode=None,
    object_metadata=None,
    preserve_ascii_control_characters=None,
    reference_file_schema_uri=None,
    file_set_spec_type=None,
    null_marker=None,
    null_markers=None,
    time_zone=None,
    date_format=None,
    datetime_format=None,
    time_format=None,
    timestamp_format=None,
    source_column_match=None,
    parquet_map_target_type=None,
    timestamp_target_precision=None,
):
  """Returns a ExternalDataConfiguration from the file or specification string.

  Determines if the input string is a file path or a string,
  then returns either the parsed file contents, or the parsed configuration from
  string. The file content is expected to be JSON representation of
  ExternalDataConfiguration. The specification is expected to be of the form
  schema@format=uri i.e. schema is separated from format and uri by '@'. If the
  uri itself contains '@' or '=' then the JSON file option should be used.
  "format=" can be omitted for CSV files.

  Raises:
    UsageError: when incorrect usage or invalid args are used.
  """
  maybe_filepath = os.path.expanduser(file_path_or_simple_spec)
  if os.path.isfile(maybe_filepath):
    try:
      with open(maybe_filepath) as external_config_file:
        return yaml.safe_load(external_config_file)
    except yaml.error.YAMLError as e:
      raise app.UsageError(
          'Error decoding YAML external table definition from file %s: %s'
          % (maybe_filepath, e)
      )
  else:
    source_format = 'CSV'
    schema = None
    connection_id = None
    error_msg = (
        'Error decoding external_table_definition. '
        'external_table_definition should either be the name of a '
        'JSON file or the text representation of an external table '
        'definition. Given:%s'
    ) % (file_path_or_simple_spec)

    parts = file_path_or_simple_spec.split('@')
    if len(parts) == 1:
      # Schema and connection are not specified.
      format_and_uri = parts[0]
    elif len(parts) == 2:
      # when there are 2 components, it can be:
      # 1. format=uri@connection_id.e.g csv=gs://bucket/file@us.conn1
      # 2. schema@format=uri        e.g.col1::INTEGER@csv=gs://bucket/file
      # if the first element is format=uri, then second element is connnection.
      # Else, the first is schema, second is format=uri.
      if parts[0].find('://') >= 0:
        # format=uri and connection specified.
        format_and_uri = parts[0]
        connection_id = parts[1]
      else:
        # Schema and format=uri are specified.
        schema = parts[0]
        format_and_uri = parts[1]
    elif len(parts) == 3:
      # Schema and connection both are specified
      schema = parts[0]
      format_and_uri = parts[1]
      connection_id = parts[2]
    else:
      raise app.UsageError(error_msg)

    separator_pos = format_and_uri.find('=')
    if separator_pos < 0:
      # Format is not specified
      uri = format_and_uri
    else:
      source_format = format_and_uri[0:separator_pos]
      uri = format_and_uri[separator_pos + 1 :]

    if not uri:
      raise app.UsageError(error_msg)
    # When using short notation for external table definition
    # autodetect is always performed.

    return CreateExternalTableDefinition(
        source_format,
        uri,
        schema,
        True,
        connection_id,
        use_avro_logical_types=use_avro_logical_types,
        parquet_enum_as_string=parquet_enum_as_string,
        parquet_enable_list_inference=parquet_enable_list_inference,
        metadata_cache_mode=metadata_cache_mode,
        object_metadata=object_metadata,
        preserve_ascii_control_characters=preserve_ascii_control_characters,
        reference_file_schema_uri=reference_file_schema_uri,
        file_set_spec_type=file_set_spec_type,
        null_marker=null_marker,
        null_markers=null_markers,
        time_zone=time_zone,
        date_format=date_format,
        datetime_format=datetime_format,
        time_format=time_format,
        timestamp_format=timestamp_format,
        source_column_match=source_column_match,
        parquet_map_target_type=parquet_map_target_type,
        timestamp_target_precision=timestamp_target_precision,
    )


def GetJson(
    file_path_or_json_string: str,
) -> Optional[Dict[str, Any]]:
  """Returns a JSON object from the file or a JSON string.

  Determines if the input string is a file path or a string,
  then returns either the parsed file contents, or the parsed JSON from
  string. The file content is expected to be a JSON string.

  Args:
    file_path_or_json_string: Path to the JSON file or a JSON string.

  Raises:
    UsageError: when incorrect usage or invalid args are used.
  """
  maybe_filepath = os.path.expanduser(file_path_or_json_string)
  if os.path.isfile(maybe_filepath):
    try:
      with open(maybe_filepath) as json_file:
        return json.load(json_file)
    except json.decoder.JSONDecodeError as e:
      raise app.UsageError(
          'Error decoding JSON from file %s: %s' % (maybe_filepath, e)
      )
  else:
    try:
      return json.loads(file_path_or_json_string)
    except json.decoder.JSONDecodeError as e:
      raise app.UsageError(
          'Error decoding JSON from string %s: %s'
          % (file_path_or_json_string, e)
      )


def UpdateExternalCatalogTableOptions(
    current_options: Dict[str, Any],
    external_options_str: str,
) -> Dict[str, Any]:
  """Updates the external catalog table options.

  Args:
    current_options: The current external catalog table options.
    external_options_str: The new external catalog table options as a JSON
      string or a file path.

  Returns:
    The updated external catalog table options.
  """
  # Clear the parameters if they are present in the existing options but not
  # in the new external catalog table options.
  if current_options.get(_PARAMETERS_KEY) is not None:
    current_options[_PARAMETERS_KEY] = {
        k: None for k in current_options[_PARAMETERS_KEY]
    }
  # Clear the storage descriptor if they are present in the existing options
  # but not in the new external catalog table options.
  if current_options.get(_STORAGE_DESCRIPTOR_KEY) is not None:
    current_options[_STORAGE_DESCRIPTOR_KEY] = {
        k: None for k in current_options[_STORAGE_DESCRIPTOR_KEY]
    }

  external_catalog_table_options_dict = GetJson(external_options_str)
  if _PARAMETERS_KEY in external_catalog_table_options_dict:
    current_options.setdefault(_PARAMETERS_KEY, {})
    current_options[_PARAMETERS_KEY].update(
        external_catalog_table_options_dict[_PARAMETERS_KEY]
    )
  else:
    current_options[_PARAMETERS_KEY] = None
  if _STORAGE_DESCRIPTOR_KEY in external_catalog_table_options_dict:
    current_options[_STORAGE_DESCRIPTOR_KEY] = (
        external_catalog_table_options_dict[_STORAGE_DESCRIPTOR_KEY]
    )
  else:
    current_options[_STORAGE_DESCRIPTOR_KEY] = None
  if _CONNECTION_ID_KEY in external_catalog_table_options_dict:
    current_options[_CONNECTION_ID_KEY] = external_catalog_table_options_dict[
        _CONNECTION_ID_KEY
    ]
  else:
    current_options[_CONNECTION_ID_KEY] = None
  return current_options


def UpdateExternalCatalogDatasetOptions(
    current_options: Dict[str, Any],
    external_options_str: str,
) -> Dict[str, Any]:
  """Updates the external catalog dataset options.

  Args:
    current_options: The current external catalog dataset options.
    external_options_str: The new external catalog dataset options as a JSON
      string or a file path.

  Returns:
    The updated external catalog dataset options.
  """
  # Clear the parameters if they are present in the existing dataset but not
  # in the new external catalog dataset options.
  if current_options.get(_PARAMETERS_KEY) is not None:
    current_options[_PARAMETERS_KEY] = {
        k: None for k in current_options[_PARAMETERS_KEY]
    }
  external_catalog_dataset_options_dict = GetJson(external_options_str)
  if _PARAMETERS_KEY in external_catalog_dataset_options_dict:
    current_options.setdefault(_PARAMETERS_KEY, {})
    for key, value in external_catalog_dataset_options_dict[
        _PARAMETERS_KEY
    ].items():
      current_options[_PARAMETERS_KEY][key] = value
  else:
    current_options[_PARAMETERS_KEY] = None
  if _DEFAULT_STORAGE_LOCATION_URI_KEY in external_catalog_dataset_options_dict:
    current_options[_DEFAULT_STORAGE_LOCATION_URI_KEY] = (
        external_catalog_dataset_options_dict[_DEFAULT_STORAGE_LOCATION_URI_KEY]
    )
  else:
    current_options[_DEFAULT_STORAGE_LOCATION_URI_KEY] = None
  return current_options


def PrintPageToken(page_token):
  """Prints the page token in the pretty format.

  Args:
    page_token: The dictionary mapping of pageToken with string 'nextPageToken'.
  """
  formatter = utils_flags.get_formatter_from_flags(secondary_format='pretty')
  utils_formatting.configure_formatter(
      formatter, bq_id_utils.ApiClientHelper.NextPageTokenReference
  )
  formatter.AddDict(page_token)
  formatter.Print()


def ParseTimePartitioning(
    partitioning_type=None,
    partitioning_expiration=None,
    partitioning_field=None,
    partitioning_minimum_partition_date=None,
    partitioning_require_partition_filter=None,
):
  """Parses time partitioning from the arguments.

  Args:
    partitioning_type: type for the time partitioning. Supported types are HOUR,
      DAY, MONTH, and YEAR. The default value is DAY when other arguments are
      specified, which generates one partition per day.
    partitioning_expiration: number of seconds to keep the storage for a
      partition. A negative value clears this setting.
    partitioning_field: if not set, the table is partitioned based on the
      loading time; if set, the table is partitioned based on the value of this
      field.
    partitioning_minimum_partition_date: lower boundary of partition date for
      field based partitioning table.
    partitioning_require_partition_filter: if true, queries on the table must
      have a partition filter so not all partitions are scanned.

  Returns:
    Time partitioning if any of the arguments is not None, otherwise None.

  Raises:
    UsageError: when failed to parse.
  """

  time_partitioning = {}
  key_type = 'type'
  key_expiration = 'expirationMs'
  key_field = 'field'
  key_minimum_partition_date = 'minimumPartitionDate'
  key_require_partition_filter = 'requirePartitionFilter'
  if partitioning_type is not None:
    time_partitioning[key_type] = partitioning_type
  if partitioning_expiration is not None:
    time_partitioning[key_expiration] = partitioning_expiration * 1000
  if partitioning_field is not None:
    time_partitioning[key_field] = partitioning_field
  if partitioning_minimum_partition_date is not None:
    if partitioning_field is not None:
      time_partitioning[key_minimum_partition_date] = (
          partitioning_minimum_partition_date
      )
    else:
      raise app.UsageError(
          'Need to specify --time_partitioning_field for '
          '--time_partitioning_minimum_partition_date.'
      )
  if partitioning_require_partition_filter is not None:
    if time_partitioning:
      time_partitioning[key_require_partition_filter] = (
          partitioning_require_partition_filter
      )

  if time_partitioning:
    if key_type not in time_partitioning:
      time_partitioning[key_type] = 'DAY'
    if (
        key_expiration in time_partitioning
        and time_partitioning[key_expiration] <= 0
    ):
      time_partitioning[key_expiration] = None
    return time_partitioning
  else:
    return None


def ParseFileSetSpecType(file_set_spec_type=None):
  """Parses the file set specification type from the arguments.

  Args:
    file_set_spec_type: specifies how to discover files given source URIs.

  Returns:
    file set specification type.
  Raises:
    UsageError: when an illegal value is passed.
  """
  if file_set_spec_type is None:
    return None
  valid_spec_types = ['FILE_SYSTEM_MATCH', 'NEW_LINE_DELIMITED_MANIFEST']
  if file_set_spec_type not in valid_spec_types:
    raise app.UsageError(
        'Error parsing file_set_spec_type, only FILE_SYSTEM_MATCH, '
        'NEW_LINE_DELIMITED_MANIFEST or no value are accepted'
    )
  return 'FILE_SET_SPEC_TYPE_' + file_set_spec_type


def ParseClustering(
    clustering_fields: Optional[str] = None,
) -> Optional[Dict[str, List[str]]]:
  """Parses clustering from the arguments.

  Args:
    clustering_fields: Comma-separated field names.

  Returns:
    Clustering if any of the arguments is not None, otherwise None. Special
    case if clustering_fields is passed in as an empty string instead of None,
    in which case we'll return {}, to support the scenario where user wants to
    update a table and remove the clustering spec.
  """

  if clustering_fields == '':  # pylint: disable=g-explicit-bool-comparison
    return {}
  elif clustering_fields is not None:
    return {'fields': clustering_fields.split(',')}
  else:
    return None


def ParseNumericTypeConversionMode(
    numeric_type_conversion_mode: Optional[str] = None,
) -> Optional[str]:
  """Parses the numeric type conversion mode from the arguments.

  Args:
    numeric_type_conversion_mode: specifies how the numeric values are handled
      when the value is out of scale.

  Returns:
    The conversion mode.

  Raises:
    UsageError: when an illegal value is passed.
  """

  if numeric_type_conversion_mode is None:
    return None
  elif numeric_type_conversion_mode == 'ROUND':
    return 'NUMERIC_TYPE_VALUE_ROUND'
  else:
    raise app.UsageError(
        'Error parsing numeric_type_conversion_mode, only ROUND or no value '
        'are accepted'
    )


def ParseRangePartitioning(range_partitioning_spec=None):
  """Parses range partitioning from the arguments.

  Args:
    range_partitioning_spec: specification for range partitioning in the format
      of field,start,end,interval.

  Returns:
    Range partitioning if range_partitioning_spec is not None, otherwise None.
  Raises:
    UsageError: when the spec fails to parse.
  """

  range_partitioning = {}
  key_field = 'field'
  key_range = 'range'
  key_range_start = 'start'
  key_range_end = 'end'
  key_range_interval = 'interval'

  if range_partitioning_spec is not None:
    parts = range_partitioning_spec.split(',')
    if len(parts) != 4:
      raise app.UsageError(
          'Error parsing range_partitioning. range_partitioning should be in '
          'the format of "field,start,end,interval"'
      )
    range_partitioning[key_field] = parts[0]
    range_spec = {}
    range_spec[key_range_start] = parts[1]
    range_spec[key_range_end] = parts[2]
    range_spec[key_range_interval] = parts[3]
    range_partitioning[key_range] = range_spec

  if range_partitioning:
    return range_partitioning
  else:
    return None


def IsSuccessfulDmlOrDdlJob(printable_job_info: str) -> bool:
  """Returns True iff the job is successful and is a DML/DDL query job."""
  return (
      'Affected Rows' in printable_job_info
      or 'DDL Operation Performed' in printable_job_info
  )


def MaybeGetSessionTempObjectName(
    dataset_id: str, object_id: str
) -> Optional[str]:
  """If we have a session temporary object, returns the user name of the object.

  Args:
    dataset_id: Dataset of object
    object_id: Id of object

  Returns:
    If the object is a session temp object, the name of the object after
    stripping out internal stuff such as session prefix and signature encodings.

    If the object is not a session temp object, the return value is None.
  """
  if not re.fullmatch('_[0-9a-f]{40}', dataset_id):
    return None  # Not an anonymous dataset

  session_prefix_regexp = (
      '_[0-9a-f]{8}_[0-9a-f]{4}_[0-9a-f]{4}_[0-9a-f]{4}_[0-9a-f]{12}_'
  )
  opt_signature_encoding_regexp = '(?:_b0a98f6_.*)?'
  match = re.fullmatch(
      session_prefix_regexp + '(.*?)' + opt_signature_encoding_regexp, object_id
  )
  if not match:
    return None  # No session prefix
  return match.group(1)


def PrintJobMessages(printable_job_info):
  """Prints additional info from a job formatted for printing.

  If the job had a fatal error, non-fatal warnings are not shown.

  If any error/warning does not have a 'message' key, printable_job_info must
  have 'jobReference' identifying the job.

  For DML queries prints number of affected rows.
  For DDL queries prints the performed operation and the target.
  """
  messages = GetJobMessagesForPrinting(printable_job_info)
  if messages:
    print(messages)


def GetJobMessagesForPrinting(printable_job_info):
  """Similar to _PrintJobMessages(), but returns a string, rather than printing."""
  result_lines = []

  job_ref = '(unknown)'  # Should never be seen, but beats a weird crash.
  if 'jobReference' in printable_job_info:
    job_ref = printable_job_info['jobReference']

  # For failing jobs, display the error but not any warnings, because those
  # may be more distracting than helpful.
  if printable_job_info['State'] == 'FAILURE':
    error_result = printable_job_info['status']['errorResult']
    error_ls = printable_job_info['status'].get('errors', [])
    error = bq_error.CreateBigqueryError(error_result, error_result, error_ls)
    result_lines.append(
        'Error encountered during job execution:\n%s\n' % (error,)
    )
  elif 'errors' in printable_job_info['status']:
    warnings = printable_job_info['status']['errors']
    result_lines.append((
        'Warning%s encountered during job execution:\n'
        % ('' if len(warnings) == 1 else 's')
    ))
    recommend_show = False
    for w in warnings:
      # Some warnings include detailed error messages, and some just
      # include programmatic error codes.  Some have a 'location'
      # separately, and some put it in the 'message' text.
      if 'message' not in w:
        recommend_show = True
      else:
        if 'location' in w:
          message = '[%s] %s' % (w['location'], w['message'])
        else:
          message = w['message']
        if message is not None:
          message = message.encode('utf-8')
        result_lines.append('%s\n' % message)
    if recommend_show:
      result_lines.append('Use "bq show -j %s" to view job warnings.' % job_ref)
  elif 'Affected Rows' in printable_job_info:
    result_lines.append(
        'Number of affected rows: %s\n' % printable_job_info['Affected Rows']
    )
  elif 'DDL Target Table' in printable_job_info:
    ddl_target_table = printable_job_info['DDL Target Table']
    project_id = ddl_target_table.get('projectId')
    dataset_id = ddl_target_table.get('datasetId')
    table_id = ddl_target_table.get('tableId')
    op = _DDL_OPERATION_MAP.get(
        printable_job_info.get('DDL Operation Performed')
    )
    # DDL Target Table is returned for both TABLE DDL and DROP ALL ROW ACCESS
    # POLICIES DDL statements.
    if project_id and dataset_id and table_id and op:
      if 'DDL Affected Row Access Policy Count' in printable_job_info:
        ddl_affected_row_access_policy_count = printable_job_info[
            'DDL Affected Row Access Policy Count'
        ]
        result_lines.append(
            '{op} {count} row access policies on table '
            '{project}.{dataset}.{table}\n'.format(
                op=op,
                count=ddl_affected_row_access_policy_count,
                project=project_id,
                dataset=dataset_id,
                table=table_id,
            )
        )
      elif (
          'Statement Type' in printable_job_info
          and 'INDEX' in printable_job_info['Statement Type']
      ):
        if 'SEARCH_INDEX' in printable_job_info['Statement Type']:
          result_lines.append(
              '%s search index on table %s.%s.%s\n'
              % (
                  stringutil.ensure_str(op),
                  stringutil.ensure_str(project_id),
                  stringutil.ensure_str(dataset_id),
                  stringutil.ensure_str(table_id),
              )
          )
        elif 'VECTOR_INDEX' in printable_job_info['Statement Type']:
          index_progress_instruction = ''
          if printable_job_info.get('DDL Operation Performed') in (
              'CREATE',
              'REPLACE',
          ):
            index_progress_instruction = (
                'Please query %s.%s.INFORMATION_SCHEMA to check the progress '
                ' of the index.\n'
            ) % (project_id, dataset_id)
          result_lines.append(
              '%s vector index on table %s.%s.%s\n%s'
              % (
                  stringutil.ensure_str(op),
                  stringutil.ensure_str(project_id),
                  stringutil.ensure_str(dataset_id),
                  stringutil.ensure_str(table_id),
                  stringutil.ensure_str(index_progress_instruction),
              )
          )
      else:
        result_lines.append(
            '%s %s.%s.%s\n'
            % (
                stringutil.ensure_str(op),
                stringutil.ensure_str(project_id),
                stringutil.ensure_str(dataset_id),
                stringutil.ensure_str(table_id),
            )
        )
        if 'Default Connection Stats' in printable_job_info:
          default_connection_stats = printable_job_info[
              'Default Connection Stats'
          ]
          location_id = job_ref['location']
          if 'provisioned' in default_connection_stats:
            if printable_job_info['Statement Type'] == 'CREATE_MODEL':
              target_type = 'model'
            else:
              target_type = 'table'
            result_lines.append(
                'Default connection created for %s [%s] in project [%s] in'
                ' region [%s]\n'
                % (
                    stringutil.ensure_str(target_type),
                    stringutil.ensure_str(table_id),
                    stringutil.ensure_str(project_id),
                    stringutil.ensure_str(location_id),
                )
            )
          if 'permissionUpdated' in default_connection_stats:
            result_lines.append(
                'Your IAM policy has been updated for the default connection\n'
            )
  elif 'DDL Target Routine' in printable_job_info:
    ddl_target_routine = printable_job_info['DDL Target Routine']
    project_id = ddl_target_routine.get('projectId')
    dataset_id = ddl_target_routine.get('datasetId')
    routine_id = ddl_target_routine.get('routineId')
    op = _DDL_OPERATION_MAP.get(
        printable_job_info.get('DDL Operation Performed')
    )
    temp_object_name = MaybeGetSessionTempObjectName(dataset_id, routine_id)
    if temp_object_name is not None:
      result_lines.append('%s temporary routine %s' % (op, temp_object_name))
    else:
      result_lines.append(
          '%s %s.%s.%s' % (op, project_id, dataset_id, routine_id)
      )
  elif 'DDL Target Row Access Policy' in printable_job_info:
    ddl_target_row_access_policy = printable_job_info[
        'DDL Target Row Access Policy'
    ]
    project_id = ddl_target_row_access_policy.get('projectId')
    dataset_id = ddl_target_row_access_policy.get('datasetId')
    table_id = ddl_target_row_access_policy.get('tableId')
    row_access_policy_id = ddl_target_row_access_policy.get('policyId')
    op = _DDL_OPERATION_MAP.get(
        printable_job_info.get('DDL Operation Performed')
    )
    if project_id and dataset_id and table_id and row_access_policy_id and op:
      result_lines.append(
          '{op} row access policy {policy} on table {project}.{dataset}.{table}'
          .format(
              op=op,
              policy=row_access_policy_id,
              project=project_id,
              dataset=dataset_id,
              table=table_id,
          )
      )
  elif 'Assertion' in printable_job_info:
    result_lines.append('Assertion successful')

  if 'Session Id' in printable_job_info:
    result_lines.append('In session: %s' % printable_job_info['Session Id'])

  return '\n'.join(result_lines)


def PrintObjectInfo(
    object_info,
    reference: bq_id_utils.ApiClientHelper.Reference,
    custom_format: bq_consts.CustomPrintFormat,
    print_reference: bool = True,
) -> None:
  """Prints the object with various formats."""
  # The JSON formats are handled separately so that they don't print
  # the record as a list of one record.
  if custom_format == 'schema':
    if 'schema' not in object_info or 'fields' not in object_info['schema']:
      raise app.UsageError('Unable to retrieve schema from specified table.')
    bq_utils.PrintFormattedJsonObject(object_info['schema']['fields'])
  elif FLAGS.format in ['prettyjson', 'json']:
    bq_utils.PrintFormattedJsonObject(object_info)
  elif FLAGS.format in [None, 'sparse', 'pretty']:
    formatter = utils_flags.get_formatter_from_flags()
    utils_formatting.configure_formatter(
        formatter,
        type(reference),
        print_format=custom_format,
        object_info=object_info,
    )
    object_info = utils_formatting.format_info_by_type(
        object_info, type(reference)
    )
    if object_info:
      formatter.AddDict(object_info)
    if reference.typename and print_reference:
      print('%s %s\n' % (reference.typename.capitalize(), reference))
    formatter.Print()
    print()
    if isinstance(reference, bq_id_utils.ApiClientHelper.JobReference):
      PrintJobMessages(object_info)
  else:
    formatter = utils_flags.get_formatter_from_flags()
    formatter.AddColumns(list(object_info.keys()))
    formatter.AddDict(object_info)
    formatter.Print()


def PrintObjectsArray(object_infos, objects_type):
  if FLAGS.format in ['prettyjson', 'json']:
    bq_utils.PrintFormattedJsonObject(object_infos)
  elif FLAGS.format in [None, 'sparse', 'pretty']:
    if not object_infos:
      return
    formatter = utils_flags.get_formatter_from_flags()
    utils_formatting.configure_formatter(
        formatter, objects_type, print_format='list'
    )
    formatted_infos = list(
        map(
            functools.partial(
                utils_formatting.format_info_by_type,
                object_type=objects_type,
            ),
            object_infos,
        )
    )
    for info in formatted_infos:
      formatter.AddDict(info)
    formatter.Print()
  elif object_infos:
    formatter = utils_flags.get_formatter_from_flags()
    formatter.AddColumns(list(object_infos[0].keys()))
    for info in object_infos:
      formatter.AddDict(info)
    formatter.Print()


class ResourceMetadata(TypedDict):
  token: NotRequired[str] = None
  unreachable: NotRequired[List[str]] = None


def PrintObjectsArrayWithMetadata(
    objects_list: List[Any],
    objects_type: Type[bq_id_utils.ApiClientHelper.Reference],
    passed_flags: NamedTuple(
        'PassedFlags',
        [
            ('print_last_token', bool),
            ('print_unreachable', bool),
        ],
    ),
    objects_metadata: Optional[ResourceMetadata],
) -> None:
  """Prints the objects array with metadata configured to print using flags.

  If there is no `objects_metadata` passed in, then this function has the same
  behaviour as `PrintObjectsArray`.

  Different metadata can be printed by setting flags in `passed_flags`. With
  a `format` of 'sparse' or 'pretty' then nothing will be printed if no flags
  are set. With a format of 'prettyjson' or 'json' then the `objects_list`
  will be printed as a `results` value even if no flags are set to print
  metadata but if some are set, these will also be printed.

  Arguments:
    objects_list: The list of resources to print.
    objects_type: The type of the resources to be printed.
    passed_flags: Flags used to configure the printing behaviour.
    objects_metadata: Optional metadata to be printed.
  """
  if FLAGS.format in ['prettyjson', 'json']:
    if passed_flags.print_last_token or passed_flags.print_unreachable:
      json_object = {'results': objects_list}
      if passed_flags.print_last_token and 'token' in objects_metadata:
        json_object['token'] = objects_metadata['token']
      if passed_flags.print_unreachable and 'unreachable' in objects_metadata:
        json_object['unreachable'] = objects_metadata['unreachable']
    else:
      json_object = objects_list
    bq_utils.PrintFormattedJsonObject(json_object)
  elif FLAGS.format in [None, 'sparse', 'pretty']:
    PrintObjectsArray(objects_list, objects_type)
    if objects_metadata is None:
      return
    if passed_flags.print_last_token and 'token' in objects_metadata:
      print('\nNext token: ' + objects_metadata['token'])
    if passed_flags.print_unreachable and 'unreachable' in objects_metadata:
      print('\nUnreachable: ' + ', '.join(objects_metadata['unreachable']))


def ParseUdfResources(udf_resources):
  """Parses UDF resources from an array of resource URIs.

  Arguments:
    udf_resources: Array of udf resource URIs.

  Returns:
    Array of UDF resources parsed into the format expected by the BigQuery API
    client.
  """

  if udf_resources is None:
    return None
  inline_udf_resources = []
  external_udf_resources = []
  for uris in udf_resources:
    for uri in uris.split(','):
      if os.path.isfile(uri):
        with open(uri) as udf_file:
          inline_udf_resources.append(udf_file.read())
      else:
        if not uri.startswith('gs://'):
          raise app.UsageError(
              'Non-inline resources must be Google Cloud Storage (gs://) URIs'
          )
        external_udf_resources.append(uri)
  udfs = []
  if inline_udf_resources:
    for udf_code in inline_udf_resources:
      udfs.append({'inlineCode': udf_code})
  if external_udf_resources:
    for uri in external_udf_resources:
      udfs.append({'resourceUri': uri})
  return udfs


def ValidateDatasetName(dataset_name: str) -> None:
  """A regex to ensure the dataset name is valid.


  Arguments:
    dataset_name: string name of the dataset to be validated.

  Raises:
    UsageError: An error occurred due to invalid dataset string.
  """
  is_valid = re.fullmatch(r'[a-zA-Z0-9\_]{1,1024}', dataset_name)
  if not is_valid:
    raise app.UsageError(
        'Dataset name: %s is invalid, must be letters '
        '(uppercase or lowercase), numbers, and underscores up to '
        '1024 characters.' % dataset_name
    )


def ParseParameters(parameters):
  """Parses query parameters from an array of name:type:value.

  Arguments:
    parameters: An iterable of string-form query parameters: name:type:value.
      Name may be omitted to indicate a positional parameter: :type:value. Type
      may be omitted to indicate a string: name::value, or ::value.

  Returns:
    A list of query parameters in the form for the BigQuery API client.
  """
  if not parameters:
    return None
  results = []
  for param_string in parameters:
    if os.path.isfile(param_string):
      with open(param_string) as f:
        results += json.load(f)
    else:
      results.append(ParseParameter(param_string))
  return results


def SplitParam(param_string):
  split = param_string.split(':', 1)
  if len(split) != 2:
    raise app.UsageError(
        'Query parameters must be of the form: '
        '"name:type:value", ":type:value", or "name::value". '
        'An empty name produces a positional parameter. '
        'An empty type produces a STRING parameter.'
    )
  return split


def ParseParameter(param_string):
  """Parse a string of the form <name><type>:<value> into each part."""
  name, param_string = SplitParam(param_string)
  try:
    type_dict, value_dict = ParseParameterTypeAndValue(param_string)
  except app.UsageError as e:
    print('Error parsing parameter %s: %s' % (name, e))
    sys.exit(1)
  result = {'parameterType': type_dict, 'parameterValue': value_dict}
  if name:
    result['name'] = name
  return result


def ParseParameterTypeAndValue(param_string):
  """Parse a string of the form <recursive_type>:<value> into each part."""
  type_string, value_string = SplitParam(param_string)
  if not type_string:
    type_string = 'STRING'
  type_dict = ParseParameterType(type_string)
  return type_dict, ParseParameterValue(type_dict, value_string)


def ParseParameterType(type_string):
  """Parse a parameter type string into a JSON dict for the BigQuery API."""
  type_dict = {'type': type_string.upper()}
  if type_string.upper().startswith('ARRAY<') and type_string.endswith('>'):
    type_dict = {
        'type': 'ARRAY',
        'arrayType': ParseParameterType(type_string[6:-1]),
    }
  if type_string.startswith('STRUCT<') and type_string.endswith('>'):
    type_dict = {
        'type': 'STRUCT',
        'structTypes': ParseStructType(type_string[7:-1]),
    }
  if type_string.startswith('RANGE<') and type_string.endswith('>'):
    type_dict = {
        'type': 'RANGE',
        'rangeElementType': ParseParameterType(type_string[6:-1]),
    }
  if not type_string:
    raise app.UsageError('Query parameter missing type')
  return type_dict


def ParseStructType(type_string):
  """Parse a Struct QueryParameter type into a JSON dict form."""
  subtypes = []
  for name, sub_type in StructTypeSplit(type_string):
    subtypes.append({'type': ParseParameterType(sub_type), 'name': name})
  return subtypes


def StructTypeSplit(type_string):
  """Yields single field-name, sub-types tuple from a StructType string.

  Raises:
    UsageError: When a field name is missing.
  """
  while type_string:
    next_span = type_string.split(',', 1)[0]
    if '<' in next_span:
      angle_count = 0
      i = 0
      for i in range(next_span.find('<'), len(type_string)):
        if type_string[i] == '<':
          angle_count += 1
        if type_string[i] == '>':
          angle_count -= 1
        if angle_count == 0:
          break
      if angle_count != 0:
        raise app.UsageError('Malformatted struct type')
      next_span = type_string[: i + 1]
    type_string = type_string[len(next_span) + 1 :]
    splits = next_span.split(None, 1)
    if len(splits) != 2:
      raise app.UsageError('Struct parameter missing name for field')
    yield splits


def FormatRfc3339(datetime_obj: datetime.datetime) -> str:
  """Formats a datetime.datetime object (UTC) in RFC3339.

  https://developers.google.com/protocol-buffers/docs/reference/google.protobuf#timestamp

  Args:
    datetime_obj: A datetime.datetime object representing a datetime in UTC.

  Returns:
    The string representation of the date in RFC3339.
  """
  return datetime_obj.isoformat('T') + 'Z'


def ParseRangeParameterValue(range_value: str) -> Tuple[str, str]:
  """Parse a range parameter value string into its components.

  Args:
    range_value: A range value string of the form "[<start>, <end>)".

  Returns:
    A tuple (<start>, <end>).

  Raises:
    app.UsageError: if the input range value string was not formatted correctly.
  """
  parsed = ParseRangeString(range_value)
  if parsed is None:
    raise app.UsageError(
        f'Invalid range parameter value: {range_value}. Expected format:'
        ' "[<start>, <end>)"'
    )
  return parsed


def ParseParameterValue(type_dict, value_input):
  """Parse a parameter value of type `type_dict` from value_input.

  Arguments:
    type_dict: The JSON-dict type as which to parse `value_input`.
    value_input: Either a string representing the value, or a JSON dict for
      array and value types.
  """
  if 'structTypes' in type_dict:
    if isinstance(value_input, str):
      if value_input == 'NULL':
        return {'structValues': None}
      value_input = json.loads(value_input)
    type_map = dict([(x['name'], x['type']) for x in type_dict['structTypes']])
    values = {}
    for field_name, value in value_input.items():
      values[field_name] = ParseParameterValue(type_map[field_name], value)
    return {'structValues': values}
  if 'arrayType' in type_dict:
    if isinstance(value_input, str):
      if value_input == 'NULL':
        return {'arrayValues': None}
      try:
        value_input = json.loads(value_input)
      except json.decoder.JSONDecodeError:
        tb = sys.exc_info()[2]
        # pylint: disable=raise-missing-from
        raise app.UsageError(
            'Error parsing string as JSON: %s' % value_input
        ).with_traceback(tb)
    values = [
        ParseParameterValue(type_dict['arrayType'], x) for x in value_input
    ]
    if not values:  # Workaround to pass empty array parameter.
      return {'value': {}}  # An empty arrayValues list is the same as NULL.
    return {'arrayValues': values}
  if 'rangeElementType' in type_dict:
    if value_input == 'NULL':
      return {'rangeValue': None}
    start, end = ParseRangeParameterValue(value_input)
    return {
        'rangeValue': {
            'start': ParseParameterValue(type_dict['rangeElementType'], start),
            'end': ParseParameterValue(type_dict['rangeElementType'], end),
        }
    }
  return {'value': value_input if value_input != 'NULL' else None}