<commit_before>import re

import pandas as pd

import ibis
import ibis.expr.types as ir
import ibis.expr.datatypes as dt
from ibis.client import Database, Query, SQLClient
from ibis.bigquery import compiler as comp
import google.cloud.bigquery
from google.api.core.exceptions import BadRequest


NATIVE_PARTITION_COL = '_PARTITIONTIME'


def _ensure_split(table_id, dataset_id):
    split = table_id.split('.')
    if len(split) > 1:
        assert len(split) == 2
        if dataset_id:
            raise ValueError(
                "Can't pass a fully qualified table name *AND* a dataset_id"
            )
        (dataset_id, table_id) = split
    return (table_id, dataset_id)


class BigQueryCursor(object):
    """Cursor to allow the BigQuery client to reuse machinery in ibis/client.py
    """

    def __init__(self, query):
        self.query = query

    def fetchall(self):
        return list(self.query.fetch_data())

    @property
    def columns(self):
        return [field.name for field in self.query.schema]

    def __enter__(self):
        # For compatibility when constructed from Query.execute()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass


class BigQuery(Query):

    def _fetch(self, cursor):
        return pd.DataFrame(cursor.fetchall(), columns=cursor.columns)


class BigQueryAPIProxy(object):

    def __init__(self, project_id):
        self._client = google.cloud.bigquery.Client(project_id)

    @property
    def client(self):
        return self._client

    @property
    def project_id(self):
        return self.client.project

    def get_datasets(self):
        return list(self.client.list_datasets())

    def get_dataset(self, dataset_id):
        return self.client.dataset(dataset_id)

    def get_table(self, table_id, dataset_id, reload=True):
        (table_id, dataset_id) = _ensure_split(table_id, dataset_id)
        table = self.client.dataset(dataset_id).table(table_id)
        if reload:
            table.reload()
        return table

    def get_schema(self, table_id, dataset_id):
        return self.get_table(table_id, dataset_id).schema


class BigQueryDatabase(Database):
    pass


class BigQueryClient(SQLClient):

    sync_query = BigQuery
    database_class = BigQueryDatabase
    proxy_class = BigQueryAPIProxy
    dialect = comp.BigQueryDialect

    def __init__(self, project_id, dataset_id):
        self._proxy = self.__class__.proxy_class(project_id)
        self._dataset_id = dataset_id

    @property
    def project_id(self):
        return self._proxy.project_id

    @property
    def dataset_id(self):
        return self._dataset_id

    @property
    def _table_expr_klass(self):
        return ir.TableExpr

    def table(self, *args, **kwargs):
        t = super(BigQueryClient, self).table(*args, **kwargs)
        if NATIVE_PARTITION_COL in t.columns:
            col = ibis.options.bigquery.partition_col
            assert col not in t
            return (t
                    .mutate(**{col: t[NATIVE_PARTITION_COL]})
                    .drop([NATIVE_PARTITION_COL]))
        return t

    def _build_ast(self, expr, params=None):
        return comp.build_ast(expr, params=params)

    def _fully_qualified_name(self, name, database):
        dataset_id = database or self.dataset_id
        return dataset_id + '.' + name

    def _get_table_schema(self, qualified_name):
        return self.get_schema(qualified_name)

    def _execute(self, stmt, results=True):
        # TODO(phillipc): Allow **kwargs in calls to execute
        query = self._proxy.client.run_sync_query(stmt)
        query.use_legacy_sql = False
        query.run()
        return BigQueryCursor(query)

    def database(self, name=None):<commit_msg>Ensure that BigQuery queries run until they are complete, instead of until a default timeout<commit_after>import re
import time

import pandas as pd

import ibis
import ibis.expr.types as ir
import ibis.expr.datatypes as dt
from ibis.client import Database, Query, SQLClient
from ibis.bigquery import compiler as comp
import google.cloud.bigquery
from google.api.core.exceptions import BadRequest


NATIVE_PARTITION_COL = '_PARTITIONTIME'


def _ensure_split(table_id, dataset_id):
    split = table_id.split('.')
    if len(split) > 1:
        assert len(split) == 2
        if dataset_id:
            raise ValueError(
                "Can't pass a fully qualified table name *AND* a dataset_id"
            )
        (dataset_id, table_id) = split
    return (table_id, dataset_id)


class BigQueryCursor(object):
    """Cursor to allow the BigQuery client to reuse machinery in ibis/client.py
    """

    def __init__(self, query):
        self.query = query

    def fetchall(self):
        return list(self.query.fetch_data())

    @property
    def columns(self):
        return [field.name for field in self.query.schema]

    def __enter__(self):
        # For compatibility when constructed from Query.execute()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass


class BigQuery(Query):

    def _fetch(self, cursor):
        return pd.DataFrame(cursor.fetchall(), columns=cursor.columns)


class BigQueryAPIProxy(object):

    def __init__(self, project_id):
        self._client = google.cloud.bigquery.Client(project_id)

    @property
    def client(self):
        return self._client

    @property
    def project_id(self):
        return self.client.project

    def get_datasets(self):
        return list(self.client.list_datasets())

    def get_dataset(self, dataset_id):
        return self.client.dataset(dataset_id)

    def get_table(self, table_id, dataset_id, reload=True):
        (table_id, dataset_id) = _ensure_split(table_id, dataset_id)
        table = self.client.dataset(dataset_id).table(table_id)
        if reload:
            table.reload()
        return table

    def get_schema(self, table_id, dataset_id):
        return self.get_table(table_id, dataset_id).schema

    def run_sync_query(self, stmt):
        query = self.client.run_sync_query(stmt)
        query.use_legacy_sql = False
        query.run()
        # run_sync_query is not really synchronous: there's a timeout
        while not query.job.done():
            query.job.reload()
            time.sleep(1)
        return query


class BigQueryDatabase(Database):
    pass


class BigQueryClient(SQLClient):

    sync_query = BigQuery
    database_class = BigQueryDatabase
    proxy_class = BigQueryAPIProxy
    dialect = comp.BigQueryDialect

    def __init__(self, project_id, dataset_id):
        self._proxy = self.__class__.proxy_class(project_id)
        self._dataset_id = dataset_id

    @property
    def project_id(self):
        return self._proxy.project_id

    @property
    def dataset_id(self):
        return self._dataset_id

    @property
    def _table_expr_klass(self):
        return ir.TableExpr

    def table(self, *args, **kwargs):
        t = super(BigQueryClient, self).table(*args, **kwargs)
        if NATIVE_PARTITION_COL in t.columns:
            col = ibis.options.bigquery.partition_col
            assert col not in t
            return (t
                    .mutate(**{col: t[NATIVE_PARTITION_COL]})
                    .drop([NATIVE_PARTITION_COL]))
        return t

    def _build_ast(self, expr, params=None):
        return comp.build_ast(expr, params=params)

    def _fully_qualified_name(self, name, database):
        dataset_id = database or self.dataset_id
        return dataset_id + '.' + name

    def _get_table_schema(self, qualified_name):
        return self.get_schema(qualified_name)

    def _execute(self, stmt, results=True):
        # TODO(phillipc): Allow **kwargs in calls to execute
        query = self._proxy.run_sync_query(stmt)
        return BigQueryCursor(query)

    def database(self, name=None):
{'fraction_non_alphanumeric': 0.08857868020304568, 'fraction_numerical': 0.0006345177664974619, 'mean_word_length': 2.7780441035474595, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, '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': '6576803', 'n_tokens_mistral': 2401, 'n_tokens_neox': 2314, 'n_words': 589}