<commit_before>import logging
import contextlib
# from threading import get_ident
from time import time

from flask import current_app
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.orm.exc import DetachedInstanceError

from .repr import ReprMixin


logger = logging.getLogger('flask-vgavro-utils.sqla')


class SQLAlchemy(SQLAlchemy):
    """
    Extends driver and session with custom options.
    """
    def __init__(self, app=None, use_native_unicode=True, session_options={}, **kwargs):
        if app and 'SQLALCHEMY_AUTOCOMMIT' in app.config:
            session_options['autocommit'] = app.config['SQLALCHEMY_AUTOCOMMIT']
        super().__init__(app, use_native_unicode, session_options, **kwargs)

    def apply_driver_hacks(self, app, info, options):
        super().apply_driver_hacks(app, info, options)
        if (app and 'SQLALCHEMY_STATEMENT_TIMEOUT' in app.config and
           info.drivername in ('postgresql', 'postgresql+psycopg2')):
            timeout = int(app.config['SQLALCHEMY_STATEMENT_TIMEOUT']) * 1000
            assert 'connect_args' not in options
            options['connect_args'] = {'options': '-c statement_timeout={:d}'.format(timeout)}


class ModelReprMixin(ReprMixin):
    # See for all fields https://stackoverflow.com/a/2448930/450103
    def __repr__(self, *args, **kwargs):
        try:
            return super().__repr__(*args, **kwargs)
        except DetachedInstanceError as exc:
            return '<{} {!r}>'.format(self.__class__.__name__, exc)

    def to_dict(self, *args, **kwargs):
        fields = args or [c.name for c in self.__table__.columns]
        return super().to_dict(*fields, **kwargs)


class InstantDefaultsMixin:
    # https://github.com/kvesteri/sqlalchemy-utils/blob/master/sqlalchemy_utils/listeners.py#L24
    # instant_defaults_listener
    # TODO: maybe remove it and use from sqlalchemy_utils?
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        for key, column in self.__table__.columns.items():
            if (
                hasattr(column, 'default') and
                column.default is not None and
                getattr(self, key, None) is None and
                key not in kwargs
            ):
                if callable(column.default.arg):
                    setattr(self, key, column.default.arg(self))
                else:
                    setattr(self, key, column.default.arg)


def db_reinit(db=None, bind=None):
    """Reinitialize database"""
    from sqlalchemy.schema import DropTable
    from sqlalchemy.ext.compiler import compiles

    @compiles(DropTable, "postgresql")
    def _compile_drop_table(element, compiler, **kwargs):
        return compiler.visit_drop_table(element) + " CASCADE"

    if not db:
        db = current_app.extensions['sqlalchemy'].db
    # NOTE: bind=None to drop_all and create_all only default database
    db.drop_all(bind=bind)
    db.create_all(bind=bind)
    db.session.commit()


class transaction(contextlib.ContextDecorator):
    def __init__(self, commit=False, rollback=False, session=None,
                 ctx_name=None, logger=None):
        assert not (commit and rollback), 'Specify commit or rollback, not both'
        self.session, self.commit, self.rollback, self.ctx_name, self.logger = \
            session, commit, rollback, ctx_name, logger

    def __call__(self, func):
        if not self.ctx_name:
            self.ctx_name = func.__name__
        return super().__call__(func)

    def __enter__(self):
        if not self.session:
            self.session = current_app.extensions['sqlalchemy'].db.session
<commit_msg>sqla: allow anonymous funcs with transaction<commit_after>import logging
import contextlib
# from threading import get_ident
from time import time

from flask import current_app
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.orm.exc import DetachedInstanceError

from .repr import ReprMixin


logger = logging.getLogger('flask-vgavro-utils.sqla')


class SQLAlchemy(SQLAlchemy):
    """
    Extends driver and session with custom options.
    """
    def __init__(self, app=None, use_native_unicode=True, session_options={}, **kwargs):
        if app and 'SQLALCHEMY_AUTOCOMMIT' in app.config:
            session_options['autocommit'] = app.config['SQLALCHEMY_AUTOCOMMIT']
        super().__init__(app, use_native_unicode, session_options, **kwargs)

    def apply_driver_hacks(self, app, info, options):
        super().apply_driver_hacks(app, info, options)
        if (app and 'SQLALCHEMY_STATEMENT_TIMEOUT' in app.config and
           info.drivername in ('postgresql', 'postgresql+psycopg2')):
            timeout = int(app.config['SQLALCHEMY_STATEMENT_TIMEOUT']) * 1000
            assert 'connect_args' not in options
            options['connect_args'] = {'options': '-c statement_timeout={:d}'.format(timeout)}


class ModelReprMixin(ReprMixin):
    # See for all fields https://stackoverflow.com/a/2448930/450103
    def __repr__(self, *args, **kwargs):
        try:
            return super().__repr__(*args, **kwargs)
        except DetachedInstanceError as exc:
            return '<{} {!r}>'.format(self.__class__.__name__, exc)

    def to_dict(self, *args, **kwargs):
        fields = args or [c.name for c in self.__table__.columns]
        return super().to_dict(*fields, **kwargs)


class InstantDefaultsMixin:
    # https://github.com/kvesteri/sqlalchemy-utils/blob/master/sqlalchemy_utils/listeners.py#L24
    # instant_defaults_listener
    # TODO: maybe remove it and use from sqlalchemy_utils?
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        for key, column in self.__table__.columns.items():
            if (
                hasattr(column, 'default') and
                column.default is not None and
                getattr(self, key, None) is None and
                key not in kwargs
            ):
                if callable(column.default.arg):
                    setattr(self, key, column.default.arg(self))
                else:
                    setattr(self, key, column.default.arg)


def db_reinit(db=None, bind=None):
    """Reinitialize database"""
    from sqlalchemy.schema import DropTable
    from sqlalchemy.ext.compiler import compiles

    @compiles(DropTable, "postgresql")
    def _compile_drop_table(element, compiler, **kwargs):
        return compiler.visit_drop_table(element) + " CASCADE"

    if not db:
        db = current_app.extensions['sqlalchemy'].db
    # NOTE: bind=None to drop_all and create_all only default database
    db.drop_all(bind=bind)
    db.create_all(bind=bind)
    db.session.commit()


class transaction(contextlib.ContextDecorator):
    def __init__(self, commit=False, rollback=False, session=None,
                 ctx_name=None, logger=None):
        assert not (commit and rollback), 'Specify commit or rollback, not both'
        self.session, self.commit, self.rollback, self.ctx_name, self.logger = \
            session, commit, rollback, ctx_name, logger

    def __call__(self, func):
        if not self.ctx_name:
            self.ctx_name = getattr(func, '__name__', 'unknown')
        return super().__call__(func)

    def __enter__(self):
        if not self.session:
            self.session = current_app.extensions['sqlalchemy'].db.session

{'fraction_non_alphanumeric': 0.106703146374829, 'fraction_numerical': 0.005471956224350205, 'mean_word_length': 2.9798584648884052, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 4, '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': '4792616', 'n_tokens_mistral': 2248, 'n_tokens_neox': 2143, 'n_words': 527}