<commit_before>from datetime import datetime
import json
import string


from sqlalchemy import Column, Integer, String, DateTime, Boolean
from sqlalchemy.sql import select, func

from twisted.internet.defer import inlineCallbacks, returnValue

from aludel.database import (
    PrefixedTableCollection, make_table, TableMissingError,
)


class UniqueCodeError(Exception):
    pass


class AuditMismatch(UniqueCodeError):
    pass


class NoUniqueCodePool(UniqueCodeError):
    pass


class CannotRedeemUniqueCode(UniqueCodeError):
    def __init__(self, reason, unique_code):
        super(CannotRedeemUniqueCode, self).__init__(reason)
        self.reason = reason
        self.unique_code = unique_code


class UniqueCodePool(PrefixedTableCollection):
    # We assume all unique codes match this.
    UNIQUE_CODE_ALLOWED_CHARS = string.lowercase + string.digits

    unique_codes = make_table(
        Column("id", Integer(), primary_key=True),
        Column("unique_code", String(), nullable=False),
        Column("flavour", String()),
        Column("used", Boolean(), default=False),
        Column("created_at", DateTime(timezone=True)),
        Column("modified_at", DateTime(timezone=True)),
        Column("reason", String(), default=None),
    )

    audit = make_table(
        Column("id", Integer(), primary_key=True),
        Column("request_id", String(), nullable=False, index=True,
               unique=True),
        Column("transaction_id", String(), nullable=False, index=True),
        Column("user_id", String(), nullable=False, index=True),
        Column("request_data", String(), nullable=False),
        Column("response_data", String(), nullable=False),
        Column("error", Boolean(), nullable=False),
        Column("created_at", DateTime(timezone=True)),
        Column("unique_code", String(), index=True),
    )

    import_audit = make_table(
        Column("id", Integer(), primary_key=True),
        Column("request_id", String(), nullable=False, index=True,
               unique=True),
        Column("content_md5", String(), nullable=False),
        Column("created_at", DateTime(timezone=True)),
    )

    @inlineCallbacks
    def execute_query(self, query, *args, **kw):
        try:
            result = yield super(UniqueCodePool, self).execute_query(
                query, *args, **kw)
        except TableMissingError:
            raise NoUniqueCodePool(self.name)
        returnValue(result)

    @classmethod
    def canonicalise_unique_code(cls, unique_code):
        """Clean spurious characters out of unique code entries.

        The unique_code is cleaned by converting it to lowercase and stripping
        out any characters not in :attr:`UNIQUE_CODE_ALLOWED_CHARS`.
        """
        return ''.join(c for c in unique_code.lower()
                       if c in cls.UNIQUE_CODE_ALLOWED_CHARS)

    def _audit_request(self, audit_params, req_data, resp_data, unique_code,
                       error=False):
        request_id = audit_params['request_id']
        transaction_id = audit_params['transaction_id']
        user_id = audit_params['user_id']
        return self.execute_query(
            self.audit.insert().values(**{
                'request_id': request_id,
                'transaction_id': transaction_id,
                'user_id': user_id,
                'request_data': json.dumps(req_data),
                'response_data': json.dumps(resp_data),
                'error': error,
                'created_at': datetime.utcnow(),
                'unique_code': unique_code,
            }))

    @inlineCallbacks
    def _get_previous_request(self, audit_params, req_data):
        rows = yield self.execute_fetchall(
            self.audit.select().where(
                self.audit.c.request_id == audit_params['request_id']))
        if not rows:
            returnValue(None)
        [row] = rows
        old_audit_params = {
            'request_id': row['request_id'],
            'transaction_id': row['transaction_id'],
            'user_id': row['user_id'],
        }
        old_req_data = json.loads(row['request_data'])
        old_resp_data = json.loads(row['response_data'])
        if audit_params != old_audit_params or req_data != old_req_data:
            raise AuditMismatch()

        if row['error']:
            raise CannotRedeemUniqueCode(
                old_resp_data['reason'], old_resp_data['unique_code'])
        returnValue(json.loads(row['response_data']))

    @inlineCallbacks
    def import_unique_codes(self, request_id, content_md5, unique_code_dicts):
        # We may not have a table if this is the first import.
        yield self.create_tables()

        trx = yield self._conn.begin()

        # Check if we've already done this one.
        rows = yield self.execute_fetchall(
            self.import_audit.select().where(
                self.import_audit.c.request_id == request_id))
        if rows:
            yield trx.rollback()
            [row] = rows
            if row['content_md5'] == content_md5:
                returnValue(None)
            else:
                raise AuditMismatch(row['content_md5'])

        yield self.execute_query(
            self.import_audit.insert().values(**{
                'request_id': request_id,
                'content_md5': content_md5,
                'created_at': datetime.utcnow(),
            }))

        # NOTE: We're assuming that this will be fast enough. If it isn't,
        # we'll need to make a database-specific plan of some kind.
        now = datetime.utcnow()
        unique_code_rows = [{
            'flavour': unique_code_dict['flavour'],
            'unique_code': unique_code_dict['unique_code'],
            'created_at': now,
            'modified_at': now,
        } for unique_code_dict in unique_code_dicts]
        result = yield self.execute_query(
            self.unique_codes.insert(), unique_code_rows)
        yield trx.commit()
        returnValue(result)

    def _format_unique_code(self, unique_code_row, fields=None):
        if fields is None:
            fields = set(f for f in unique_code_row.keys()
                         if f not in ('created_at', 'modified_at'))
        return dict((k, v) for k, v in unique_code_row.items() if k in fields)<commit_msg>Remove unnecessary uses of ** with dict literals.<commit_after>from datetime import datetime
import json
import string


from sqlalchemy import Column, Integer, String, DateTime, Boolean
from sqlalchemy.sql import select, func

from twisted.internet.defer import inlineCallbacks, returnValue

from aludel.database import (
    PrefixedTableCollection, make_table, TableMissingError,
)


class UniqueCodeError(Exception):
    pass


class AuditMismatch(UniqueCodeError):
    pass


class NoUniqueCodePool(UniqueCodeError):
    pass


class CannotRedeemUniqueCode(UniqueCodeError):
    def __init__(self, reason, unique_code):
        super(CannotRedeemUniqueCode, self).__init__(reason)
        self.reason = reason
        self.unique_code = unique_code


class UniqueCodePool(PrefixedTableCollection):
    # We assume all unique codes match this.
    UNIQUE_CODE_ALLOWED_CHARS = string.lowercase + string.digits

    unique_codes = make_table(
        Column("id", Integer(), primary_key=True),
        Column("unique_code", String(), nullable=False),
        Column("flavour", String()),
        Column("used", Boolean(), default=False),
        Column("created_at", DateTime(timezone=True)),
        Column("modified_at", DateTime(timezone=True)),
        Column("reason", String(), default=None),
    )

    audit = make_table(
        Column("id", Integer(), primary_key=True),
        Column("request_id", String(), nullable=False, index=True,
               unique=True),
        Column("transaction_id", String(), nullable=False, index=True),
        Column("user_id", String(), nullable=False, index=True),
        Column("request_data", String(), nullable=False),
        Column("response_data", String(), nullable=False),
        Column("error", Boolean(), nullable=False),
        Column("created_at", DateTime(timezone=True)),
        Column("unique_code", String(), index=True),
    )

    import_audit = make_table(
        Column("id", Integer(), primary_key=True),
        Column("request_id", String(), nullable=False, index=True,
               unique=True),
        Column("content_md5", String(), nullable=False),
        Column("created_at", DateTime(timezone=True)),
    )

    @inlineCallbacks
    def execute_query(self, query, *args, **kw):
        try:
            result = yield super(UniqueCodePool, self).execute_query(
                query, *args, **kw)
        except TableMissingError:
            raise NoUniqueCodePool(self.name)
        returnValue(result)

    @classmethod
    def canonicalise_unique_code(cls, unique_code):
        """Clean spurious characters out of unique code entries.

        The unique_code is cleaned by converting it to lowercase and stripping
        out any characters not in :attr:`UNIQUE_CODE_ALLOWED_CHARS`.
        """
        return ''.join(c for c in unique_code.lower()
                       if c in cls.UNIQUE_CODE_ALLOWED_CHARS)

    def _audit_request(self, audit_params, req_data, resp_data, unique_code,
                       error=False):
        request_id = audit_params['request_id']
        transaction_id = audit_params['transaction_id']
        user_id = audit_params['user_id']
        return self.execute_query(
            self.audit.insert().values(
                request_id=request_id,
                transaction_id=transaction_id,
                user_id=user_id,
                request_data=json.dumps(req_data),
                response_data=json.dumps(resp_data),
                error=error,
                created_at=datetime.utcnow(),
                unique_code=unique_code,
            ))

    @inlineCallbacks
    def _get_previous_request(self, audit_params, req_data):
        rows = yield self.execute_fetchall(
            self.audit.select().where(
                self.audit.c.request_id == audit_params['request_id']))
        if not rows:
            returnValue(None)
        [row] = rows
        old_audit_params = {
            'request_id': row['request_id'],
            'transaction_id': row['transaction_id'],
            'user_id': row['user_id'],
        }
        old_req_data = json.loads(row['request_data'])
        old_resp_data = json.loads(row['response_data'])
        if audit_params != old_audit_params or req_data != old_req_data:
            raise AuditMismatch()

        if row['error']:
            raise CannotRedeemUniqueCode(
                old_resp_data['reason'], old_resp_data['unique_code'])
        returnValue(json.loads(row['response_data']))

    @inlineCallbacks
    def import_unique_codes(self, request_id, content_md5, unique_code_dicts):
        # We may not have a table if this is the first import.
        yield self.create_tables()

        trx = yield self._conn.begin()

        # Check if we've already done this one.
        rows = yield self.execute_fetchall(
            self.import_audit.select().where(
                self.import_audit.c.request_id == request_id))
        if rows:
            yield trx.rollback()
            [row] = rows
            if row['content_md5'] == content_md5:
                returnValue(None)
            else:
                raise AuditMismatch(row['content_md5'])

        yield self.execute_query(
            self.import_audit.insert().values(
                request_id=request_id,
                content_md5=content_md5,
                created_at=datetime.utcnow(),
            ))

        # NOTE: We're assuming that this will be fast enough. If it isn't,
        # we'll need to make a database-specific plan of some kind.
        now = datetime.utcnow()
        unique_code_rows = [{
            'flavour': unique_code_dict['flavour'],
            'unique_code': unique_code_dict['unique_code'],
            'created_at': now,
            'modified_at': now,
        } for unique_code_dict in unique_code_dicts]
        result = yield self.execute_query(
            self.unique_codes.insert(), unique_code_rows)
        yield trx.commit()
        returnValue(result)

    def _format_unique_code(self, unique_code_row, fields=None):
        if fields is None:
            fields = set(f for f in unique_code_row.keys()
                         if f not in ('created_at', 'modified_at'))
        return dict((k, v) for k, v in unique_code_row.items() if k in fields)
{'fraction_non_alphanumeric': 0.10674202021810077, 'fraction_numerical': 0.0011143835071240945, 'mean_word_length': 2.5232753785754345, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5424177', 'n_tokens_mistral': 3616, 'n_tokens_neox': 3510, 'n_words': 794}