<commit_before>#!/usr/bin/env python

"""Wrapper around 'bq' tool from Google Cloud SDK.

It automates some common tasks, like dataset and table creation and schema
changes.

It reads a definition of the desired state of tables from config files,
compares it to what's in BigQuery now, and performs necessary updates. It will
not execute any removals though. The user must execute removals manually, to
acknowledge the data loss.

Prerequisites:
  * 'bq' tool is in PATH.
  * Application default credentials are configured.
"""

import argparse
import ast
import collections
import difflib
import json
import logging
import subprocess
import sys
import tempfile


class Error(Exception):
  pass


class BigQueryError(Error):
  pass


class DefinitionError(Error):
  pass


class BigQuery(object):
  """A thin wrapper around 'bq' tool.

  Hides its CLI interface and argument marshaling. In theory, can be replaced
  with direct REST API calls. It will speed everything up significantly.

  All methods may raise BigQueryError.
  """

  def __init__(self, project_id):
    self.project_id = str(project_id)

  def _call(self, cmd, args, decode_json=True):
    full_cmd = [
      'bq',
      '--api_version', 'v2',  # we rely on particular format of JSON responses
      '--project_id', self.project_id,
      '--format', 'prettyjson',
      cmd
    ] + list(args)
    logging.debug('Calling %s', full_cmd)
    try:
      output = subprocess.check_output(full_cmd)
    except subprocess.CalledProcessError as exc:
      logging.debug('Failed with code %d: %s', exc.returncode)
      raise BigQueryError(exc.output.strip())
    logging.debug('Output:\n%s', output or '<empty>')
    if decode_json:
      return json.loads(output) if output else None
    return output

  def _tid(self, dataset_id, table_id):
    """Returns fully-qualified table ID."""
    return '%s:%s.%s' % (self.project_id, dataset_id, table_id)

  ### Datasets

  def list_datasets(self):
    """Returns a list of dataset IDs in the project."""
    reply = self._call('ls', ['%s:' % self.project_id]) or []
    return [d['datasetId'] for d in reply]

  def create_dataset(self, dataset_id, description=''):
    # 'mk' returns "Dataset '...' successfully created." instead of JSON.
    self._call(
      'mk', ['--description', description, '-d', dataset_id],
      decode_json=False)

  ### Tables

  def list_tables(self, dataset_id):
    """Returns a list of table IDs within the dataset."""
    reply = self._call('ls', ['%s:%s.' % (self.project_id, dataset_id)]) or []
    return [t['tableId'] for t in reply]

  def get_table_metadata(self, dataset_id, table_id):
    """Returns information about an existing table (as a dict).

    The format of the dict matches 'bigquery#table' resource defined here:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/tables
    """
    out = self._call('show', [self._tid(dataset_id, table_id)])
    assert out.get('kind') == 'bigquery#table'
    return out

  def create_table(
      self,
      dataset_id,
      table_id,
      schema,
      time_partitioning,
      time_partitioning_exp_days,
      description):
    """Creates a table (time partitioned or not).<commit_msg>Fix bq_util.py tool.<commit_after>#!/usr/bin/env python

"""Wrapper around 'bq' tool from Google Cloud SDK.

It automates some common tasks, like dataset and table creation and schema
changes.

It reads a definition of the desired state of tables from config files,
compares it to what's in BigQuery now, and performs necessary updates. It will
not execute any removals though. The user must execute removals manually, to
acknowledge the data loss.

Prerequisites:
  * 'bq' tool is in PATH.
  * Application default credentials are configured.
"""

import argparse
import ast
import collections
import difflib
import json
import logging
import subprocess
import sys
import tempfile


class Error(Exception):
  pass


class BigQueryError(Error):
  pass


class DefinitionError(Error):
  pass


class BigQuery(object):
  """A thin wrapper around 'bq' tool.

  Hides its CLI interface and argument marshaling. In theory, can be replaced
  with direct REST API calls. It will speed everything up significantly.

  All methods may raise BigQueryError.
  """

  def __init__(self, project_id):
    self.project_id = str(project_id)

  def _call(self, cmd, args, decode_json=True):
    full_cmd = [
      'bq',
      '--api_version', 'v2',  # we rely on particular format of JSON responses
      '--project_id', self.project_id,
      '--format', 'prettyjson',
      cmd
    ] + list(args)
    logging.debug('Calling %s', full_cmd)
    try:
      output = subprocess.check_output(full_cmd)
    except subprocess.CalledProcessError as exc:
      logging.debug('Failed with code %d: %s', exc.returncode)
      raise BigQueryError(exc.output.strip())
    logging.debug('Output:\n%s', output or '<empty>')
    if decode_json:
      return json.loads(output) if output else None
    return output

  def _tid(self, dataset_id, table_id):
    """Returns fully-qualified table ID."""
    return '%s:%s.%s' % (self.project_id, dataset_id, table_id)

  ### Datasets

  def list_datasets(self):
    """Returns a list of dataset IDs in the project."""
    reply = self._call('ls', ['%s:' % self.project_id]) or []
    return [d['datasetReference']['datasetId'] for d in reply]

  def create_dataset(self, dataset_id, description=''):
    # 'mk' returns "Dataset '...' successfully created." instead of JSON.
    self._call(
      'mk', ['--description', description, '-d', dataset_id],
      decode_json=False)

  ### Tables

  def list_tables(self, dataset_id):
    """Returns a list of table IDs within the dataset."""
    reply = self._call('ls', ['%s:%s.' % (self.project_id, dataset_id)]) or []
    return [t['tableReference']['tableId'] for t in reply]

  def get_table_metadata(self, dataset_id, table_id):
    """Returns information about an existing table (as a dict).

    The format of the dict matches 'bigquery#table' resource defined here:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/tables
    """
    out = self._call('show', [self._tid(dataset_id, table_id)])
    assert out.get('kind') == 'bigquery#table'
    return out

  def create_table(
      self,
      dataset_id,
      table_id,
      schema,
      time_partitioning,
      time_partitioning_exp_days,
      description):
    """Creates a table (time partitioned or not).
{'fraction_non_alphanumeric': 0.11478206917946332, 'fraction_numerical': 0.0006204436171862882, 'mean_word_length': 3.5698086463501064, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3797116', 'n_tokens_mistral': 1978, 'n_tokens_neox': 1917, 'n_words': 657}