<commit_before>from contextlib import contextmanager, ContextDecorator
from threading import local

from django.conf import settings
from django.core import signals
from django.db import DEFAULT_DB_ALIAS
from django.utils.functional import cached_property

import sqlalchemy
from memoized import memoized
from six.moves.urllib.parse import urlencode
from sqlalchemy.orm.scoping import scoped_session
from sqlalchemy.orm.session import sessionmaker

from corehq.util.test_utils import unit_testing_only

from .util import select_db_for_read

DEFAULT_ENGINE_ID = DEFAULT_DB_ALIAS
UCR_ENGINE_ID = 'ucr'
AAA_DB_ENGINE_ID = 'aaa-data'
ICDS_UCR_CITUS_ENGINE_ID = 'icds-ucr-citus'


def get_icds_ucr_citus_db_alias():
    return _get_db_alias_or_none(ICDS_UCR_CITUS_ENGINE_ID)


def get_aaa_db_alias():
    return _get_db_alias_or_none(AAA_DB_ENGINE_ID)


def _get_db_alias_or_none(enigne_id):
    try:
        return connection_manager.get_django_db_alias(enigne_id)
    except KeyError:
        return None


_thread_locals = local()


READ_FROM_CITUS_STANDBYS = 'READ_FROM_CITUS_STANDBYS'


def allow_read_from_citus_standbys():
    return getattr(_thread_locals, READ_FROM_CITUS_STANDBYS, False)


class read_from_citus_standbys(ContextDecorator):
    def __enter__(self):
        setattr(_thread_locals, READ_FROM_CITUS_STANDBYS, True)

    def __exit__(self, exc_type, exc_val, exc_tb):
        setattr(_thread_locals, READ_FROM_CITUS_STANDBYS, False)


def is_citus_db(connection):
    """
    :param connection: either a sqlalchemy connection or a Django cursor
    """
    res = connection.execute("SELECT 1 FROM pg_extension WHERE extname = 'citus'")
    if res is None:
        res = list(connection)
    return bool(list(res))


def is_citus(engine, citus_list=[]):
    # hacking mutable default arg to memoize list of citus DBs
    #   to avoid hitting DB. citus_list is never to be actually
    #   passed.
    if engine.url in citus_list:
        return True
    with engine.begin() as connection:
        if is_citus_db(connection):
            citus_list.append(engine.url)
            return True
        else:
            return False


def create_engine(connection_string):
    # paramstyle='format' allows you to use column names that include the ')' character
    # otherwise queries will sometimes be misformated/error when formatting
    # https://github.com/zzzeek/sqlalchemy/blob/ff20903/lib/sqlalchemy/dialects/postgresql/psycopg2.py#L173
    engine = sqlalchemy.create_engine(connection_string, paramstyle='format')

    if is_citus(engine) and allow_read_from_citus_standbys():
        connect_args = {'options': '-c citus.use_secondary_nodes=always'}
        return sqlalchemy.create_engine(connection_string, paramstyle='format', connect_args=connect_args)
    else:
        return engine


class SessionHelper(object):
    """<commit_msg>avoid hitting db for non-citus; no need to check if citus<commit_after>from contextlib import contextmanager, ContextDecorator
from threading import local

from django.conf import settings
from django.core import signals
from django.db import DEFAULT_DB_ALIAS
from django.utils.functional import cached_property

import sqlalchemy
from memoized import memoized
from six.moves.urllib.parse import urlencode
from sqlalchemy.orm.scoping import scoped_session
from sqlalchemy.orm.session import sessionmaker

from corehq.util.test_utils import unit_testing_only

from .util import select_db_for_read

DEFAULT_ENGINE_ID = DEFAULT_DB_ALIAS
UCR_ENGINE_ID = 'ucr'
AAA_DB_ENGINE_ID = 'aaa-data'
ICDS_UCR_CITUS_ENGINE_ID = 'icds-ucr-citus'


def get_icds_ucr_citus_db_alias():
    return _get_db_alias_or_none(ICDS_UCR_CITUS_ENGINE_ID)


def get_aaa_db_alias():
    return _get_db_alias_or_none(AAA_DB_ENGINE_ID)


def _get_db_alias_or_none(enigne_id):
    try:
        return connection_manager.get_django_db_alias(enigne_id)
    except KeyError:
        return None


_thread_locals = local()


READ_FROM_CITUS_STANDBYS = 'READ_FROM_CITUS_STANDBYS'


def allow_read_from_citus_standbys():
    return getattr(_thread_locals, READ_FROM_CITUS_STANDBYS, False)


class read_from_citus_standbys(ContextDecorator):
    def __enter__(self):
        setattr(_thread_locals, READ_FROM_CITUS_STANDBYS, True)

    def __exit__(self, exc_type, exc_val, exc_tb):
        setattr(_thread_locals, READ_FROM_CITUS_STANDBYS, False)


def is_citus_db(connection):
    """
    :param connection: either a sqlalchemy connection or a Django cursor
    """
    res = connection.execute("SELECT 1 FROM pg_extension WHERE extname = 'citus'")
    if res is None:
        res = list(connection)
    return bool(list(res))


def is_citus(engine, citus_list=[], non_citus_list=[]):
    # hacking mutable default arg to memoize list of citus DBs
    #   to avoid hitting DB. citus_list, non_citus_list
    #   is never to be actually passed.
    if engine.url in citus_list:
        return True
    if engine.url in non_citus_list:
        return False
    with engine.begin() as connection:
        if is_citus_db(connection):
            citus_list.append(engine.url)
            return True
        else:
            non_citus_list.append(engine.url)
            return False


def create_engine(connection_string):
    # paramstyle='format' allows you to use column names that include the ')' character
    # otherwise queries will sometimes be misformated/error when formatting
    # https://github.com/zzzeek/sqlalchemy/blob/ff20903/lib/sqlalchemy/dialects/postgresql/psycopg2.py#L173
    engine = sqlalchemy.create_engine(connection_string, paramstyle='format')

    if allow_read_from_citus_standbys():
        connect_args = {'options': '-c citus.use_secondary_nodes=always'}
        return sqlalchemy.create_engine(connection_string, paramstyle='format', connect_args=connect_args)
    else:
        return engine


class SessionHelper(object):
    """
{'fraction_non_alphanumeric': 0.07150135869565218, 'fraction_numerical': 0.0033967391304347825, 'mean_word_length': 4.286355475763016, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, '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': '2106390', 'n_tokens_mistral': 2031, 'n_tokens_neox': 1905, 'n_words': 443}