<commit_before>import time
import weakref
import gevent
from socket import gethostname
from urllib import quote_plus as urlquote
from sqlalchemy import create_engine, event

from inbox.sqlalchemy_ext.util import ForceStrictMode
from inbox.config import config
from inbox.util.stats import statsd_client
from nylas.logging import get_logger, find_first_app_frame_and_name
from warnings import filterwarnings
filterwarnings('ignore', message='Invalid utf8mb4 character string')
log = get_logger()


DB_POOL_SIZE = config.get_required('DB_POOL_SIZE')
# Sane default of max overflow=5 if value missing in config.
DB_POOL_MAX_OVERFLOW = config.get('DB_POOL_MAX_OVERFLOW') or 5
DB_POOL_TIMEOUT = config.get('DB_POOL_TIMEOUT') or 60


pool_tracker = weakref.WeakKeyDictionary()


# See
# https://github.com/PyMySQL/mysqlclient-python/blob/master/samples/waiter_gevent.py
def gevent_waiter(fd, hub=gevent.hub.get_hub()):
    hub.wait(hub.loop.io(fd, 1))


def build_uri(username, password, hostname, port, database_name):
    uri_template = 'mysql+mysqldb://{username}:{password}@{hostname}' \
                   ':{port}/{database_name}?charset=utf8mb4'
    return uri_template.format(username=urlquote(username),
                               password=urlquote(password),
                               hostname=urlquote(hostname),
                               port=port,
                               database_name=urlquote(database_name))


def engine(database_name, database_uri, pool_size=DB_POOL_SIZE,
           max_overflow=DB_POOL_MAX_OVERFLOW, pool_timeout=DB_POOL_TIMEOUT,
           echo=False):
    engine = create_engine(database_uri,
                           listeners=[ForceStrictMode()],
                           isolation_level='READ COMMITTED',
                           echo=False,
                           pool_size=pool_size,
                           pool_timeout=pool_timeout,
                           pool_recycle=3600,
                           max_overflow=max_overflow,
                           connect_args={'charset': 'utf8mb4',
                                         'waiter': gevent_waiter})

    @event.listens_for(engine, 'checkout')
    def receive_checkout(dbapi_connection, connection_record,
                         connection_proxy):
        '''Log checkedout and overflow when a connection is checked out'''
        hostname = gethostname().replace(".", "-")
        process_name = str(config.get("PROCESS_NAME", "unknown"))

        statsd_client.gauge(".".join(
            ["dbconn", database_name, hostname, process_name,
             "checkedout"]),
            connection_proxy._pool.checkedout())

        statsd_client.gauge(".".join(
            ["dbconn", database_name, hostname, process_name,
             "overflow"]),
            connection_proxy._pool.overflow())

        # Keep track of where and why this connection was checked out.
        log = get_logger()
        context = log._context._dict.copy()
        f, name = find_first_app_frame_and_name(ignores=['sqlalchemy',
                                                         'inbox.ignition',
                                                         'nylas.logging'])
        source = '{}:{}'.format(name, f.f_lineno)

        pool_tracker[dbapi_connection] = {
            'source': source,
            'context': context,
            'checkedout_at': time.time()
        }

    @event.listens_for(engine, 'checkin')
    def receive_checkin(dbapi_connection, connection_record):
        if dbapi_connection in pool_tracker:
            del pool_tracker[dbapi_connection]

    return engine


class EngineManager(object):

    def __init__(self, databases, users, include_disabled=False):
        self.engines = {}
        keys = set()
        schema_names = set()
        for database in databases:
            hostname = database['HOSTNAME']
            port = database['PORT']
            username = users[hostname]['USER']
            password = users[hostname]['PASSWORD']
            for shard in database['SHARDS']:
                schema_name = shard['SCHEMA_NAME']
                key = shard['ID']

                # Perform some sanity checks on the configuration.
                assert isinstance(key, int)
                assert key not in keys, \
                    'Shard key collision: key {} is repeated'.format(key)
                assert schema_name not in schema_names, \
                    'Shard name collision: {} is repeated'.format(schema_name)<commit_msg>Connect to localhost if using proxysql<commit_after>import time
import weakref
import gevent
from socket import gethostname
from urllib import quote_plus as urlquote
from sqlalchemy import create_engine, event

from inbox.sqlalchemy_ext.util import ForceStrictMode
from inbox.config import config
from inbox.util.stats import statsd_client
from nylas.logging import get_logger, find_first_app_frame_and_name
from warnings import filterwarnings
filterwarnings('ignore', message='Invalid utf8mb4 character string')
log = get_logger()


DB_POOL_SIZE = config.get_required('DB_POOL_SIZE')
# Sane default of max overflow=5 if value missing in config.
DB_POOL_MAX_OVERFLOW = config.get('DB_POOL_MAX_OVERFLOW') or 5
DB_POOL_TIMEOUT = config.get('DB_POOL_TIMEOUT') or 60


pool_tracker = weakref.WeakKeyDictionary()


# See
# https://github.com/PyMySQL/mysqlclient-python/blob/master/samples/waiter_gevent.py
def gevent_waiter(fd, hub=gevent.hub.get_hub()):
    hub.wait(hub.loop.io(fd, 1))


def build_uri(username, password, hostname, port, database_name):
    uri_template = 'mysql+mysqldb://{username}:{password}@{hostname}' \
                   ':{port}/{database_name}?charset=utf8mb4'
    return uri_template.format(username=urlquote(username),
                               password=urlquote(password),
                               hostname=urlquote(hostname),
                               port=port,
                               database_name=urlquote(database_name))


def engine(database_name, database_uri, pool_size=DB_POOL_SIZE,
           max_overflow=DB_POOL_MAX_OVERFLOW, pool_timeout=DB_POOL_TIMEOUT,
           echo=False):
    engine = create_engine(database_uri,
                           listeners=[ForceStrictMode()],
                           isolation_level='READ COMMITTED',
                           echo=False,
                           pool_size=pool_size,
                           pool_timeout=pool_timeout,
                           pool_recycle=3600,
                           max_overflow=max_overflow,
                           connect_args={'charset': 'utf8mb4',
                                         'waiter': gevent_waiter})

    @event.listens_for(engine, 'checkout')
    def receive_checkout(dbapi_connection, connection_record,
                         connection_proxy):
        '''Log checkedout and overflow when a connection is checked out'''
        hostname = gethostname().replace(".", "-")
        process_name = str(config.get("PROCESS_NAME", "unknown"))

        statsd_client.gauge(".".join(
            ["dbconn", database_name, hostname, process_name,
             "checkedout"]),
            connection_proxy._pool.checkedout())

        statsd_client.gauge(".".join(
            ["dbconn", database_name, hostname, process_name,
             "overflow"]),
            connection_proxy._pool.overflow())

        # Keep track of where and why this connection was checked out.
        log = get_logger()
        context = log._context._dict.copy()
        f, name = find_first_app_frame_and_name(ignores=['sqlalchemy',
                                                         'inbox.ignition',
                                                         'nylas.logging'])
        source = '{}:{}'.format(name, f.f_lineno)

        pool_tracker[dbapi_connection] = {
            'source': source,
            'context': context,
            'checkedout_at': time.time()
        }

    @event.listens_for(engine, 'checkin')
    def receive_checkin(dbapi_connection, connection_record):
        if dbapi_connection in pool_tracker:
            del pool_tracker[dbapi_connection]

    return engine


class EngineManager(object):

    def __init__(self, databases, users, include_disabled=False):
        self.engines = {}
        keys = set()
        schema_names = set()
        use_proxysql = config.get('USE_PROXYSQL', False)
        for database in databases:
            hostname = '127.0.0.1' if use_proxysql else database['HOSTNAME']
            port = database['PORT']
            username = users[hostname]['USER']
            password = users[hostname]['PASSWORD']
            for shard in database['SHARDS']:
                schema_name = shard['SCHEMA_NAME']
                key = shard['ID']

                # Perform some sanity checks on the configuration.
                assert isinstance(key, int)
                assert key not in keys, \
                    'Shard key collision: key {} is repeated'.format(key)
                assert schema_name not in schema_names, \
                    'Shard name collision: {} is repeated'.format(schema_name)
{'fraction_non_alphanumeric': 0.09497756866177919, 'fraction_numerical': 0.003939161833898676, 'mean_word_length': 2.164819944598338, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1308102', 'n_tokens_mistral': 2536, 'n_tokens_neox': 2461, 'n_words': 521}