<commit_before>from __future__ import print_function

from collections import defaultdict
import string
import subprocess
import tempfile
import time
import uuid

from google.cloud import bigquery as gcbq
from google.cloud.exceptions import Conflict, NotFound

import pandas as pd

from django.conf import settings

from gcutils.storage import Client as StorageClient
from gcutils.table_dumper import TableDumper


DATASETS = {
    'hscic': settings.BQ_HSCIC_DATASET,
    'measures': settings.BQ_MEASURES_DATASET,
    'tmp_eu': settings.BQ_TMP_EU_DATASET,
}


try:
    DATASETS['test'] = settings.BQ_TEST_DATASET
except AttributeError:
    pass


class Client(object):
    def __init__(self, dataset_key=None):
        self.project_name = settings.BQ_PROJECT

        # gcbq expects an environment variable called
        # GOOGLE_APPLICATION_CREDENTIALS whose value is the path of a JSON file
        # containing the credentials to access Google Cloud Services.
        self.gcbq_client = gcbq.Client(project=self.project_name)

        if dataset_key is None:
            self.dataset_name = None
            self.dataset = None
        else:
            self.dataset_name = DATASETS[dataset_key]
            self.dataset = self.gcbq_client.dataset(self.dataset_name)

    def list_jobs(self):
        return self.gcbq_client.list_jobs()

    def create_dataset(self):
        self.dataset.location = settings.BQ_LOCATION
        self.dataset.default_table_expiration_ms =\
            settings.BQ_DEFAULT_TABLE_EXPIRATION_MS
        self.dataset.create()

    def delete_dataset(self):
        for table in self.dataset.list_tables():
            table.delete()
        self.dataset.delete()

    def create_table(self, table_name, schema):
        table = self.dataset.table(table_name, schema)

        try:
            table.create()
        except NotFound as e:
            assert 'Not found: Dataset' in str(e)
            self.create_dataset()
            table.create()

        return Table(table, self.project_name)

    def get_table(self, table_name):
        table = self.dataset.table(table_name)
        return Table(table, self.project_name)

    def get_or_create_table(self, table_name, schema):
        try:
            table = self.create_table(table_name, schema)
        except Conflict:
            table = self.get_table(table_name)
        return table

    def create_storage_backed_table(self, table_name, schema, gcs_path):
        gcs_client = StorageClient()
        bucket = gcs_client.bucket()
        if bucket.get_blob(gcs_path) is None:
            raise RuntimeError('Could not find blob at {}'.format(gcs_path))

        gcs_uri = 'gs://{}/{}'.format(self.project_name, gcs_path)

        resource = {
            'tableReference': {'tableId': table_name},
            'externalDataConfiguration': {
                'csvOptions': {'skipLeadingRows': '1'},
                'sourceFormat': 'CSV',
                'sourceUris': [gcs_uri],
                'schema': {'fields': schema}
            }
        }

        path = '/projects/{}/datasets/{}/tables'.format(
            self.project_name,
            self.dataset_name
        )

        try:
            self.gcbq_client._connection.api_request(
                method='POST',
                path=path,
                data=resource
            )
        except NotFound as e:
            assert 'Not found: Dataset' in str(e)
            self.create_dataset()
            self.gcbq_client._connection.api_request(
                method='POST',
                path=path,
                data=resource
            )

        return self.get_table(table_name)

    def create_table_with_view(self, table_name, sql, legacy):
        assert '{project}' in sql
        sql = interpolate_sql(sql, project=self.project_name)
        table = self.dataset.table(table_name)
        table.view_query = sql
        table.view_use_legacy_sql = legacy
        try:
            table.create()
        except NotFound as e:
            assert 'Not found: Dataset' in str(e)
            self.create_dataset()
            table.create()
        return Table(table, self.project_name)

    def query(self, sql, legacy=False, **options):
        sql = interpolate_sql(sql)
        query = self.gcbq_client.run_sync_query(sql)
        set_options(query, options)
        query.use_legacy_sql = legacy

        query.run()

        # The call to .run() might return before results are actually ready.
        # See https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#timeoutMs
        wait_for_job(query.job)
<commit_msg>Re-raise BQ errors<commit_after>from __future__ import print_function

from collections import defaultdict
import string
import subprocess
import tempfile
import time
import uuid

from google.cloud import bigquery as gcbq
from google.cloud.exceptions import Conflict, NotFound

import pandas as pd

from django.conf import settings

from gcutils.storage import Client as StorageClient
from gcutils.table_dumper import TableDumper


DATASETS = {
    'hscic': settings.BQ_HSCIC_DATASET,
    'measures': settings.BQ_MEASURES_DATASET,
    'tmp_eu': settings.BQ_TMP_EU_DATASET,
}


try:
    DATASETS['test'] = settings.BQ_TEST_DATASET
except AttributeError:
    pass


class Client(object):
    def __init__(self, dataset_key=None):
        self.project_name = settings.BQ_PROJECT

        # gcbq expects an environment variable called
        # GOOGLE_APPLICATION_CREDENTIALS whose value is the path of a JSON file
        # containing the credentials to access Google Cloud Services.
        self.gcbq_client = gcbq.Client(project=self.project_name)

        if dataset_key is None:
            self.dataset_name = None
            self.dataset = None
        else:
            self.dataset_name = DATASETS[dataset_key]
            self.dataset = self.gcbq_client.dataset(self.dataset_name)

    def list_jobs(self):
        return self.gcbq_client.list_jobs()

    def create_dataset(self):
        self.dataset.location = settings.BQ_LOCATION
        self.dataset.default_table_expiration_ms =\
            settings.BQ_DEFAULT_TABLE_EXPIRATION_MS
        self.dataset.create()

    def delete_dataset(self):
        for table in self.dataset.list_tables():
            table.delete()
        self.dataset.delete()

    def create_table(self, table_name, schema):
        table = self.dataset.table(table_name, schema)

        try:
            table.create()
        except NotFound as e:
            if 'Not found: Dataset' not in str(e):
                raise
            self.create_dataset()
            table.create()

        return Table(table, self.project_name)

    def get_table(self, table_name):
        table = self.dataset.table(table_name)
        return Table(table, self.project_name)

    def get_or_create_table(self, table_name, schema):
        try:
            table = self.create_table(table_name, schema)
        except Conflict:
            table = self.get_table(table_name)
        return table

    def create_storage_backed_table(self, table_name, schema, gcs_path):
        gcs_client = StorageClient()
        bucket = gcs_client.bucket()
        if bucket.get_blob(gcs_path) is None:
            raise RuntimeError('Could not find blob at {}'.format(gcs_path))

        gcs_uri = 'gs://{}/{}'.format(self.project_name, gcs_path)

        resource = {
            'tableReference': {'tableId': table_name},
            'externalDataConfiguration': {
                'csvOptions': {'skipLeadingRows': '1'},
                'sourceFormat': 'CSV',
                'sourceUris': [gcs_uri],
                'schema': {'fields': schema}
            }
        }

        path = '/projects/{}/datasets/{}/tables'.format(
            self.project_name,
            self.dataset_name
        )

        try:
            self.gcbq_client._connection.api_request(
                method='POST',
                path=path,
                data=resource
            )
        except NotFound as e:
            if 'Not found: Dataset' not in str(e):
                raise
            self.create_dataset()
            self.gcbq_client._connection.api_request(
                method='POST',
                path=path,
                data=resource
            )

        return self.get_table(table_name)

    def create_table_with_view(self, table_name, sql, legacy):
        assert '{project}' in sql
        sql = interpolate_sql(sql, project=self.project_name)
        table = self.dataset.table(table_name)
        table.view_query = sql
        table.view_use_legacy_sql = legacy
        try:
            table.create()
        except NotFound as e:
            if 'Not found: Dataset' not in str(e):
                raise
            self.create_dataset()
            table.create()
        return Table(table, self.project_name)

    def query(self, sql, legacy=False, **options):
        sql = interpolate_sql(sql)
        query = self.gcbq_client.run_sync_query(sql)
        set_options(query, options)
        query.use_legacy_sql = legacy

        query.run()

        # The call to .run() might return before results are actually ready.
        # See https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#timeoutMs
        wait_for_job(query.job)

{'fraction_non_alphanumeric': 0.08854502471523748, 'fraction_numerical': 0.0004298302170642596, 'mean_word_length': 2.21596406357982, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '7468672', 'n_tokens_mistral': 2743, 'n_tokens_neox': 2657, 'n_words': 604}