<commit_before>from contextlib import contextmanager
from dataclasses import dataclass
from typing import Optional, Any, Dict

import google.auth
import google.cloud.bigquery
import google.cloud.exceptions
from google.api_core import retry, client_info
from google.oauth2 import service_account

from dbt.utils import format_bytes
from dbt.clients import agate_helper, gcloud
from dbt.exceptions import (
    FailedToConnectException, RuntimeException, DatabaseException
)
from dbt.adapters.base import BaseConnectionManager, Credentials
from dbt.logger import GLOBAL_LOGGER as logger
from dbt.version import __version__ as dbt_version

from hologram.helpers import StrEnum


BQ_QUERY_JOB_SPLIT = '-----Query Job SQL Follows-----'


class Priority(StrEnum):
    Interactive = 'interactive'
    Batch = 'batch'


class BigQueryConnectionMethod(StrEnum):
    OAUTH = 'oauth'
    SERVICE_ACCOUNT = 'service-account'
    SERVICE_ACCOUNT_JSON = 'service-account-json'


@dataclass
class BigQueryCredentials(Credentials):
    method: BigQueryConnectionMethod
    keyfile: Optional[str] = None
    keyfile_json: Optional[Dict[str, Any]] = None
    timeout_seconds: Optional[int] = 300
    location: Optional[str] = None
    priority: Optional[Priority] = None
    retries: Optional[int] = 1
    maximum_bytes_billed: Optional[int] = None
    _ALIASES = {
        'project': 'database',
        'dataset': 'schema',
    }

    @property
    def type(self):
        return 'bigquery'

    def _connection_keys(self):
        return ('method', 'database', 'schema', 'location', 'priority',
                'timeout_seconds', 'maximum_bytes_billed')


class BigQueryConnectionManager(BaseConnectionManager):
    TYPE = 'bigquery'

    SCOPE = ('https://www.googleapis.com/auth/bigquery',
             'https://www.googleapis.com/auth/cloud-platform',
             'https://www.googleapis.com/auth/drive')

    QUERY_TIMEOUT = 300
    RETRIES = 1
    DEFAULT_INITIAL_DELAY = 1.0  # Seconds
    DEFAULT_MAXIMUM_DELAY = 1.0  # Seconds

    @classmethod
    def handle_error(cls, error, message):
        error_msg = "\n".join([item['message'] for item in error.errors])
        raise DatabaseException(error_msg)

    def clear_transaction(self):
        pass

    @contextmanager
    def exception_handler(self, sql):
        try:
            yield

        except google.cloud.exceptions.BadRequest as e:
            message = "Bad request while running query"
            self.handle_error(e, message)

        except google.cloud.exceptions.Forbidden as e:
            message = "Access denied while running query"
            self.handle_error(e, message)

        except Exception as e:
            logger.debug("Unhandled error while running:\n{}".format(sql))
            logger.debug(e)
            if isinstance(e, RuntimeException):
                # during a sql query, an internal to dbt exception was raised.
                # this sounds a lot like a signal handler and probably has
                # useful information, so raise it without modification.
                raise
            exc_message = str(e)
            # the google bigquery library likes to add the query log, which we
            # don't want to log. Hopefully they never change this!
            if BQ_QUERY_JOB_SPLIT in exc_message:
                exc_message = exc_message.split(BQ_QUERY_JOB_SPLIT)[0].strip()
            raise RuntimeException(exc_message)

    def cancel_open(self) -> None:
        pass

    @classmethod
    def close(cls, connection):
        connection.state = 'closed'

        return connection

    def begin(self):
        pass

    def commit(self):
        pass

    @classmethod
    def get_bigquery_credentials(cls, profile_credentials):
        method = profile_credentials.method
        creds = service_account.Credentials

        if method == BigQueryConnectionMethod.OAUTH:
            credentials, project_id = google.auth.default(scopes=cls.SCOPE)
            return credentials

        elif method == BigQueryConnectionMethod.SERVICE_ACCOUNT:
            keyfile = profile_credentials.keyfile
            return creds.from_service_account_file(keyfile, scopes=cls.SCOPE)

        elif method == BigQueryConnectionMethod.SERVICE_ACCOUNT_JSON:
            details = profile_credentials.keyfile_json
            return creds.from_service_account_info(details, scopes=cls.SCOPE)

        error = ('Invalid `method` in profile: "{}"'.format(method))
        raise FailedToConnectException(error)

    @classmethod
    def get_bigquery_client(cls, profile_credentials):
        database = profile_credentials.database
        creds = cls.get_bigquery_credentials(profile_credentials)
        location = getattr(profile_credentials, 'location', None)

        info = client_info.ClientInfo(user_agent=f'dbt-{dbt_version}')
        return google.cloud.bigquery.Client(
            database,
            creds,
            location=location,
            client_info=info,
        )

    @classmethod
    def open(cls, connection):
        if connection.state == 'open':
            logger.debug('Connection is already open, skipping open.')
            return connection

        try:
            handle = cls.get_bigquery_client(connection.credentials)

        except google.auth.exceptions.DefaultCredentialsError:
            logger.info("Please log into GCP to continue")
            gcloud.setup_default_credentials()

            handle = cls.get_bigquery_client(connection.credentials)

        except Exception as e:
            raise
            logger.debug("Got an error when attempting to create a bigquery "
                         "client: '{}'".format(e))

            connection.handle = None
            connection.state = 'fail'

            raise FailedToConnectException(str(e))

        connection.handle = handle
        connection.state = 'open'
        return connection

    @classmethod
    def get_timeout(cls, conn):
        credentials = conn.credentials
        return credentials.timeout_seconds

    @classmethod
    def get_retries(cls, conn) -> int:
        credentials = conn.credentials
        if credentials.retries is not None:
            return credentials.retries
        else:
            return 1

    @classmethod
    def get_table_from_response(cls, resp):
        column_names = [field.name for field in resp.schema]
        return agate_helper.table_from_data_flat(resp, column_names)

    def raw_execute(self, sql, fetch=False):
        conn = self.get_thread_connection()
        client = conn.handle

        logger.debug('On {}: {}', conn.name, sql)

        job_params = {'use_legacy_sql': False}

        priority = conn.credentials.priority
        if priority == Priority.Batch:
            job_params['priority'] = google.cloud.bigquery.QueryPriority.BATCH
        else:
            job_params[
                'priority'] = google.cloud.bigquery.QueryPriority.INTERACTIVE

        maximum_bytes_billed = conn.credentials.maximum_bytes_billed
        if maximum_bytes_billed is not None and maximum_bytes_billed!=0:
            job_params['maximum_bytes_billed'] = maximum_bytes_billed

        def fn():
            return self._query_and_results(client, sql, conn, job_params)

        query_job, iterator = self._retry_and_handle(msg=sql, conn=conn, fn=fn)

        return query_job, iterator

    def execute(self, sql, auto_begin=False, fetch=None):
        sql = self._add_query_comment(sql)
        # auto_begin is ignored on bigquery, and only included for consistency
        query_job, iterator = self.raw_execute(sql, fetch=fetch)

        if fetch:
            res = self.get_table_from_response(iterator)
        else:
            res = agate_helper.empty_table()

        if query_job.statement_type == 'CREATE_VIEW':
            status = 'CREATE VIEW'

        elif query_job.statement_type == 'CREATE_TABLE_AS_SELECT':
            conn = self.get_thread_connection()
            client = conn.handle
            table = client.get_table(query_job.destination)
            status = 'CREATE TABLE ({})'.format(table.num_rows)

        elif query_job.statement_type == 'SCRIPT':<commit_msg>added whitespace around operator to pass flake8 tests<commit_after>from contextlib import contextmanager
from dataclasses import dataclass
from typing import Optional, Any, Dict

import google.auth
import google.cloud.bigquery
import google.cloud.exceptions
from google.api_core import retry, client_info
from google.oauth2 import service_account

from dbt.utils import format_bytes
from dbt.clients import agate_helper, gcloud
from dbt.exceptions import (
    FailedToConnectException, RuntimeException, DatabaseException
)
from dbt.adapters.base import BaseConnectionManager, Credentials
from dbt.logger import GLOBAL_LOGGER as logger
from dbt.version import __version__ as dbt_version

from hologram.helpers import StrEnum


BQ_QUERY_JOB_SPLIT = '-----Query Job SQL Follows-----'


class Priority(StrEnum):
    Interactive = 'interactive'
    Batch = 'batch'


class BigQueryConnectionMethod(StrEnum):
    OAUTH = 'oauth'
    SERVICE_ACCOUNT = 'service-account'
    SERVICE_ACCOUNT_JSON = 'service-account-json'


@dataclass
class BigQueryCredentials(Credentials):
    method: BigQueryConnectionMethod
    keyfile: Optional[str] = None
    keyfile_json: Optional[Dict[str, Any]] = None
    timeout_seconds: Optional[int] = 300
    location: Optional[str] = None
    priority: Optional[Priority] = None
    retries: Optional[int] = 1
    maximum_bytes_billed: Optional[int] = None
    _ALIASES = {
        'project': 'database',
        'dataset': 'schema',
    }

    @property
    def type(self):
        return 'bigquery'

    def _connection_keys(self):
        return ('method', 'database', 'schema', 'location', 'priority',
                'timeout_seconds', 'maximum_bytes_billed')


class BigQueryConnectionManager(BaseConnectionManager):
    TYPE = 'bigquery'

    SCOPE = ('https://www.googleapis.com/auth/bigquery',
             'https://www.googleapis.com/auth/cloud-platform',
             'https://www.googleapis.com/auth/drive')

    QUERY_TIMEOUT = 300
    RETRIES = 1
    DEFAULT_INITIAL_DELAY = 1.0  # Seconds
    DEFAULT_MAXIMUM_DELAY = 1.0  # Seconds

    @classmethod
    def handle_error(cls, error, message):
        error_msg = "\n".join([item['message'] for item in error.errors])
        raise DatabaseException(error_msg)

    def clear_transaction(self):
        pass

    @contextmanager
    def exception_handler(self, sql):
        try:
            yield

        except google.cloud.exceptions.BadRequest as e:
            message = "Bad request while running query"
            self.handle_error(e, message)

        except google.cloud.exceptions.Forbidden as e:
            message = "Access denied while running query"
            self.handle_error(e, message)

        except Exception as e:
            logger.debug("Unhandled error while running:\n{}".format(sql))
            logger.debug(e)
            if isinstance(e, RuntimeException):
                # during a sql query, an internal to dbt exception was raised.
                # this sounds a lot like a signal handler and probably has
                # useful information, so raise it without modification.
                raise
            exc_message = str(e)
            # the google bigquery library likes to add the query log, which we
            # don't want to log. Hopefully they never change this!
            if BQ_QUERY_JOB_SPLIT in exc_message:
                exc_message = exc_message.split(BQ_QUERY_JOB_SPLIT)[0].strip()
            raise RuntimeException(exc_message)

    def cancel_open(self) -> None:
        pass

    @classmethod
    def close(cls, connection):
        connection.state = 'closed'

        return connection

    def begin(self):
        pass

    def commit(self):
        pass

    @classmethod
    def get_bigquery_credentials(cls, profile_credentials):
        method = profile_credentials.method
        creds = service_account.Credentials

        if method == BigQueryConnectionMethod.OAUTH:
            credentials, project_id = google.auth.default(scopes=cls.SCOPE)
            return credentials

        elif method == BigQueryConnectionMethod.SERVICE_ACCOUNT:
            keyfile = profile_credentials.keyfile
            return creds.from_service_account_file(keyfile, scopes=cls.SCOPE)

        elif method == BigQueryConnectionMethod.SERVICE_ACCOUNT_JSON:
            details = profile_credentials.keyfile_json
            return creds.from_service_account_info(details, scopes=cls.SCOPE)

        error = ('Invalid `method` in profile: "{}"'.format(method))
        raise FailedToConnectException(error)

    @classmethod
    def get_bigquery_client(cls, profile_credentials):
        database = profile_credentials.database
        creds = cls.get_bigquery_credentials(profile_credentials)
        location = getattr(profile_credentials, 'location', None)

        info = client_info.ClientInfo(user_agent=f'dbt-{dbt_version}')
        return google.cloud.bigquery.Client(
            database,
            creds,
            location=location,
            client_info=info,
        )

    @classmethod
    def open(cls, connection):
        if connection.state == 'open':
            logger.debug('Connection is already open, skipping open.')
            return connection

        try:
            handle = cls.get_bigquery_client(connection.credentials)

        except google.auth.exceptions.DefaultCredentialsError:
            logger.info("Please log into GCP to continue")
            gcloud.setup_default_credentials()

            handle = cls.get_bigquery_client(connection.credentials)

        except Exception as e:
            raise
            logger.debug("Got an error when attempting to create a bigquery "
                         "client: '{}'".format(e))

            connection.handle = None
            connection.state = 'fail'

            raise FailedToConnectException(str(e))

        connection.handle = handle
        connection.state = 'open'
        return connection

    @classmethod
    def get_timeout(cls, conn):
        credentials = conn.credentials
        return credentials.timeout_seconds

    @classmethod
    def get_retries(cls, conn) -> int:
        credentials = conn.credentials
        if credentials.retries is not None:
            return credentials.retries
        else:
            return 1

    @classmethod
    def get_table_from_response(cls, resp):
        column_names = [field.name for field in resp.schema]
        return agate_helper.table_from_data_flat(resp, column_names)

    def raw_execute(self, sql, fetch=False):
        conn = self.get_thread_connection()
        client = conn.handle

        logger.debug('On {}: {}', conn.name, sql)

        job_params = {'use_legacy_sql': False}

        priority = conn.credentials.priority
        if priority == Priority.Batch:
            job_params['priority'] = google.cloud.bigquery.QueryPriority.BATCH
        else:
            job_params[
                'priority'] = google.cloud.bigquery.QueryPriority.INTERACTIVE

        maximum_bytes_billed = conn.credentials.maximum_bytes_billed
        if maximum_bytes_billed is not None and maximum_bytes_billed != 0:
            job_params['maximum_bytes_billed'] = maximum_bytes_billed

        def fn():
            return self._query_and_results(client, sql, conn, job_params)

        query_job, iterator = self._retry_and_handle(msg=sql, conn=conn, fn=fn)

        return query_job, iterator

    def execute(self, sql, auto_begin=False, fetch=None):
        sql = self._add_query_comment(sql)
        # auto_begin is ignored on bigquery, and only included for consistency
        query_job, iterator = self.raw_execute(sql, fetch=fetch)

        if fetch:
            res = self.get_table_from_response(iterator)
        else:
            res = agate_helper.empty_table()

        if query_job.statement_type == 'CREATE_VIEW':
            status = 'CREATE VIEW'

        elif query_job.statement_type == 'CREATE_TABLE_AS_SELECT':
            conn = self.get_thread_connection()
            client = conn.handle
            table = client.get_table(query_job.destination)
            status = 'CREATE TABLE ({})'.format(table.num_rows)

        elif query_job.statement_type == 'SCRIPT':
{'fraction_non_alphanumeric': 0.0836181640625, 'fraction_numerical': 0.00201416015625, 'mean_word_length': 2.601098901098901, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 6, 'lorem ipsum': 0, 'www.': 6, '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': '2364142', 'n_tokens_mistral': 4608, 'n_tokens_neox': 4450, 'n_words': 1091}