<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import

import codecs
import copy
import logging
import os
import sys
import uuid
from datetime import datetime
from enum import Enum

import googleapiclient.discovery
from future.utils import iteritems
from google.cloud import bigquery
from google.cloud.bigquery.job import (CopyJobConfig, CreateDisposition, QueryJobConfig,
                                       WriteDisposition)
from google.cloud.exceptions import NotFound
from google.oauth2 import service_account

from bqdm.model.schema import BigQuerySchemaField
from bqdm.model.table import BigQueryTable
from bqdm.util import (dump, echo, echo_dump, echo_ndiff)

_logger = logging.getLogger(__name__)
_logger.addHandler(logging.StreamHandler(sys.stdout))
_logger.setLevel(logging.INFO)


class SchemaMigrationMode(Enum):

    SELECT_INSERT = 'select_insert'
    SELECT_INSERT_BACKUP = 'select_insert_backup'
    REPLACE = 'replace'
    REPLACE_BACKUP = 'replace_backup'
    DROP_CREATE = 'drop_create'
    DROP_CREATE_BACKUP = 'drop_create_backup'


class TableAction(object):

    def __init__(self, executor, dataset_id,
                 migration_mode=None, backup_dataset_id=None, project=None,
                 credential_file=None, no_color=False, debug=False):
        self._executor = executor
        credentials = None
        if credential_file:
            credentials = service_account.Credentials.from_service_account_file(credential_file)
        self._client = bigquery.Client(project, credentials)
        self._api_client = googleapiclient.discovery.build('bigquery', 'v2',
                                                           credentials=credentials)
        self._dataset_ref = self._client.dataset(dataset_id)
        if backup_dataset_id:
            self._backup_dataset_ref = self._client.dataset(backup_dataset_id)
        else:
            self._backup_dataset_ref = self._dataset_ref
        if migration_mode:
            self._migration_mode = SchemaMigrationMode(migration_mode)
        else:
            self._migration_mode = SchemaMigrationMode.SELECT_INSERT
        self.no_color = no_color
        if debug:
            _logger.setLevel(logging.DEBUG)

    @property
    def dataset_reference(self):
        return self._dataset_ref

    @property
    def dataset(self):
        return self._client.get_dataset(self._dataset_ref)

    @property
    def exists_dataset(self):
        try:
            self._client.get_dataset(self._dataset_ref)
            return True
        except NotFound:
            return False

    @property
    def backup_dataset_reference(self):
        return self._backup_dataset_ref

    @property
    def backup_dataset(self):
        return self._client.get_dataset(self._backup_dataset_ref)

    @property
    def exists_backup_dataset(self):
        try:
            self._client.get_dataset(self._backup_dataset_ref)
            return True
        except NotFound:
            return False

    @property
    def migration_mode(self):
        return self._migration_mode

    @staticmethod
    def get_add_tables(source, target):
        table_ids = set(t.table_id for t in target) - set(s.table_id for s in source)
        results = [t for t in target if t.table_id in table_ids]
        return len(results), tuple(results)

    @staticmethod
    def get_change_tables(source, target):
        _, add_tables = TableAction.get_add_tables(source, target)
        results = (set(target) - set(add_tables)) - set(source)
        return len(results), tuple(results)

    @staticmethod
    def get_destroy_tables(source, target):
        table_ids = set(s.table_id for s in source) - set(t.table_id for t in target)
        results = [s for s in source if s.table_id in table_ids]
        return len(results), tuple(results)

    def migrate(self, source_table, target_table, prefix='    ', fg='yellow'):
        if self._migration_mode in [SchemaMigrationMode.SELECT_INSERT_BACKUP,
                                    SchemaMigrationMode.REPLACE_BACKUP,
                                    SchemaMigrationMode.DROP_CREATE_BACKUP]:
            self.backup(source_table.table_id)

        if self._migration_mode in [SchemaMigrationMode.SELECT_INSERT,
                                    SchemaMigrationMode.SELECT_INSERT_BACKUP]:
            query_field = TableAction.build_query_field(source_table.schema, target_table.schema)
            self.select_insert(target_table.table_id, target_table.table_id, query_field)
        elif self._migration_mode in [SchemaMigrationMode.REPLACE,
                                      SchemaMigrationMode.REPLACE_BACKUP]:
            tmp_table = self.create_temporary_table(target_table)
            query_field = TableAction.build_query_field(source_table.schema, target_table.schema)
            self.select_insert(source_table.table_id, tmp_table.table_id, query_field)
            self._destroy(target_table, prefix, fg)
            self._add(target_table, prefix, fg)
            self.select_insert(tmp_table.table_id, target_table.table_id, '*')
            self._destroy(tmp_table, prefix, fg)
        elif self._migration_mode in [SchemaMigrationMode.DROP_CREATE,
                                      SchemaMigrationMode.DROP_CREATE_BACKUP]:
            self._destroy(target_table, prefix, fg)
            self._add(target_table, prefix, fg)
        else:
            raise ValueError('Unknown migration mode.')

    def backup(self, source_table_id, prefix='    ', fg='yellow'):
        source_table = self.dataset.table(source_table_id)
        backup_table_id = 'backup_{source_table_id}_{timestamp}'.format(
            source_table_id=source_table_id,
            timestamp=datetime.utcnow().strftime('%Y%m%d%H%M%S%f'))
        backup_table = self.backup_dataset.table(backup_table_id)
        job_config = CopyJobConfig()
        job_config.create_disposition = CreateDisposition.CREATE_IF_NEEDED
        job = self._client.copy_table(source_table, backup_table, job_config=job_config)
        echo('Backing up... {0}'.format(job.job_id),
             prefix=prefix, fg=fg, no_color=self.no_color)
        job.result()
        assert job.state == 'DONE'
        error_result = job.error_result
        if error_result:
            raise RuntimeError(job.errors)<commit_msg>Fix wild card column<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import

import codecs
import copy
import logging
import os
import sys
import uuid
from datetime import datetime
from enum import Enum

import googleapiclient.discovery
from future.utils import iteritems
from google.cloud import bigquery
from google.cloud.bigquery.job import (CopyJobConfig, CreateDisposition, QueryJobConfig,
                                       WriteDisposition)
from google.cloud.exceptions import NotFound
from google.oauth2 import service_account

from bqdm.model.schema import BigQuerySchemaField
from bqdm.model.table import BigQueryTable
from bqdm.util import (dump, echo, echo_dump, echo_ndiff)

_logger = logging.getLogger(__name__)
_logger.addHandler(logging.StreamHandler(sys.stdout))
_logger.setLevel(logging.INFO)


class SchemaMigrationMode(Enum):

    SELECT_INSERT = 'select_insert'
    SELECT_INSERT_BACKUP = 'select_insert_backup'
    REPLACE = 'replace'
    REPLACE_BACKUP = 'replace_backup'
    DROP_CREATE = 'drop_create'
    DROP_CREATE_BACKUP = 'drop_create_backup'


class TableAction(object):

    def __init__(self, executor, dataset_id,
                 migration_mode=None, backup_dataset_id=None, project=None,
                 credential_file=None, no_color=False, debug=False):
        self._executor = executor
        credentials = None
        if credential_file:
            credentials = service_account.Credentials.from_service_account_file(credential_file)
        self._client = bigquery.Client(project, credentials)
        self._api_client = googleapiclient.discovery.build('bigquery', 'v2',
                                                           credentials=credentials)
        self._dataset_ref = self._client.dataset(dataset_id)
        if backup_dataset_id:
            self._backup_dataset_ref = self._client.dataset(backup_dataset_id)
        else:
            self._backup_dataset_ref = self._dataset_ref
        if migration_mode:
            self._migration_mode = SchemaMigrationMode(migration_mode)
        else:
            self._migration_mode = SchemaMigrationMode.SELECT_INSERT
        self.no_color = no_color
        if debug:
            _logger.setLevel(logging.DEBUG)

    @property
    def dataset_reference(self):
        return self._dataset_ref

    @property
    def dataset(self):
        return self._client.get_dataset(self._dataset_ref)

    @property
    def exists_dataset(self):
        try:
            self._client.get_dataset(self._dataset_ref)
            return True
        except NotFound:
            return False

    @property
    def backup_dataset_reference(self):
        return self._backup_dataset_ref

    @property
    def backup_dataset(self):
        return self._client.get_dataset(self._backup_dataset_ref)

    @property
    def exists_backup_dataset(self):
        try:
            self._client.get_dataset(self._backup_dataset_ref)
            return True
        except NotFound:
            return False

    @property
    def migration_mode(self):
        return self._migration_mode

    @staticmethod
    def get_add_tables(source, target):
        table_ids = set(t.table_id for t in target) - set(s.table_id for s in source)
        results = [t for t in target if t.table_id in table_ids]
        return len(results), tuple(results)

    @staticmethod
    def get_change_tables(source, target):
        _, add_tables = TableAction.get_add_tables(source, target)
        results = (set(target) - set(add_tables)) - set(source)
        return len(results), tuple(results)

    @staticmethod
    def get_destroy_tables(source, target):
        table_ids = set(s.table_id for s in source) - set(t.table_id for t in target)
        results = [s for s in source if s.table_id in table_ids]
        return len(results), tuple(results)

    def migrate(self, source_table, target_table, prefix='    ', fg='yellow'):
        if self._migration_mode in [SchemaMigrationMode.SELECT_INSERT_BACKUP,
                                    SchemaMigrationMode.REPLACE_BACKUP,
                                    SchemaMigrationMode.DROP_CREATE_BACKUP]:
            self.backup(source_table.table_id)

        if self._migration_mode in [SchemaMigrationMode.SELECT_INSERT,
                                    SchemaMigrationMode.SELECT_INSERT_BACKUP]:
            query_field = TableAction.build_query_field(source_table.schema, target_table.schema)
            self.select_insert(target_table.table_id, target_table.table_id, query_field)
        elif self._migration_mode in [SchemaMigrationMode.REPLACE,
                                      SchemaMigrationMode.REPLACE_BACKUP]:
            tmp_table = self.create_temporary_table(target_table)
            query_field = TableAction.build_query_field(source_table.schema, target_table.schema)
            self.select_insert(source_table.table_id, tmp_table.table_id, query_field)
            self._destroy(target_table, prefix, fg)
            self._add(target_table, prefix, fg)
            query_field = TableAction.build_query_field(target_table.schema, target_table.schema)
            self.select_insert(tmp_table.table_id, target_table.table_id, query_field)
            self._destroy(tmp_table, prefix, fg)
        elif self._migration_mode in [SchemaMigrationMode.DROP_CREATE,
                                      SchemaMigrationMode.DROP_CREATE_BACKUP]:
            self._destroy(target_table, prefix, fg)
            self._add(target_table, prefix, fg)
        else:
            raise ValueError('Unknown migration mode.')

    def backup(self, source_table_id, prefix='    ', fg='yellow'):
        source_table = self.dataset.table(source_table_id)
        backup_table_id = 'backup_{source_table_id}_{timestamp}'.format(
            source_table_id=source_table_id,
            timestamp=datetime.utcnow().strftime('%Y%m%d%H%M%S%f'))
        backup_table = self.backup_dataset.table(backup_table_id)
        job_config = CopyJobConfig()
        job_config.create_disposition = CreateDisposition.CREATE_IF_NEEDED
        job = self._client.copy_table(source_table, backup_table, job_config=job_config)
        echo('Backing up... {0}'.format(job.job_id),
             prefix=prefix, fg=fg, no_color=self.no_color)
        job.result()
        assert job.state == 'DONE'
        error_result = job.error_result
        if error_result:
            raise RuntimeError(job.errors)
{'fraction_non_alphanumeric': 0.07814579907603163, 'fraction_numerical': 0.0006264192310703939, 'mean_word_length': 2.7225298746721074, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6232759', 'n_tokens_mistral': 3626, 'n_tokens_neox': 3601, 'n_words': 708}