<commit_before>"""
This module houses the main classes you will interact with,
:class:`.Cluster` and :class:`.Session`.
"""

from futures import ThreadPoolExecutor
import logging
import time
from threading import RLock, Thread, Event
import traceback
import Queue
import weakref
from functools import partial
from itertools import groupby

from cassandra import ConsistencyLevel, AuthenticationFailed
from cassandra.connection import ConnectionException
from cassandra.decoder import (QueryMessage, ResultMessage,
                               ErrorMessage, ReadTimeoutErrorMessage,
                               WriteTimeoutErrorMessage,
                               UnavailableErrorMessage,
                               OverloadedErrorMessage,
                               PrepareMessage, ExecuteMessage,
                               PreparedQueryNotFound,
                               IsBootstrappingErrorMessage, named_tuple_factory,
                               dict_factory)
from cassandra.io.asyncorereactor import AsyncoreConnection
from cassandra.metadata import Metadata
from cassandra.policies import (RoundRobinPolicy, SimpleConvictionPolicy,
                                ExponentialReconnectionPolicy, HostDistance,
                                RetryPolicy)
from cassandra.query import SimpleStatement, PreparedStatement, BoundStatement, bind_params
from cassandra.pool import (_ReconnectionHandler, _HostReconnectionHandler,
                            HostConnectionPool)


log = logging.getLogger(__name__)


DEFAULT_MIN_REQUESTS = 5
DEFAULT_MAX_REQUESTS = 100

DEFAULT_MIN_CONNECTIONS_PER_LOCAL_HOST = 2
DEFAULT_MAX_CONNECTIONS_PER_LOCAL_HOST = 8

DEFAULT_MIN_CONNECTIONS_PER_REMOTE_HOST = 1
DEFAULT_MAX_CONNECTIONS_PER_REMOTE_HOST = 2


class NoHostAvailable(Exception):
    """
    Raised when an operation is attempted but all connections are
    busy, defunct, closed, or resulted in errors when used.
    """

    errors = None
    """
    A map of the form ``{ip: exception}`` which details the particular
    Exception that was caught for each host the operation was attempted
    against.
    """

    def __init__(self, message, errors):
        Exception.__init__(self, message, errors)
        self.errors = errors


class Cluster(object):
    """
    The main class to use when interacting with a Cassandra cluster.
    Typically, one instance of this class will be created for each
    separate Cassandra cluster that your application interacts with.

    Example usage::

        >>> from cassandra.cluster import Cluster
        >>> cluster = Cluster(['192.168.1.1', '192.168.1.2'])
        >>> session = cluster.connect()
        >>> session.execute("CREATE KEYSPACE ...")
        >>> ...
        >>> cluster.shutdown()

    """

    port = 9042
    """
    The server-side port to open connections to. Defaults to 9042.
    """

    compression = True
    """
    Whether or not compression should be enabled when possible. Defaults to
    :const:`True` and attempts to use snappy compression.
    """

    auth_provider = None
    """
    An optional function that accepts one argument, the IP address of a node,
    and returns a dict of credentials for that node.
    """

    load_balancing_policy = RoundRobinPolicy()
    """
    An instance of :class:`.policies.LoadBalancingPolicy` or
    one of its subclasses.  Defaults to :class:`~.RoundRobinPolicy`.
    """

    reconnection_policy = ExponentialReconnectionPolicy(1.0, 600.0)
    """
    An instance of :class:`.policies.ReconnectionPolicy`. Defaults to an instance
    of :class:`.ExponentialReconnectionPolicy` with a base delay of one second and
    a max delay of ten minutes.
    """

    default_retry_policy = RetryPolicy()
    """
    A default :class:`.policies.RetryPolicy` instance to use for all
    :class:`.Query` objects which do not have a :attr:`~.Query.retry_policy`
    explicitly set.
    """

    conviction_policy_factory = SimpleConvictionPolicy
    """
    A factory function which creates instances of
    :class:`.policies.ConvictionPolicy`.  Defaults to
    :class:`.policies.SimpleConvictionPolicy`.
    """

    metrics_enabled = False
    """
    Whether or not metric collection is enabled.
    """

    sockopts = None
    """
    An optional list of tuples which will be used as arguments to
    ``socket.setsockopt()`` for all created sockets.
    """

    max_schema_agreement_wait = 10
    """
    The maximum duration (in seconds) that the driver will wait for schema
    agreement across the cluster. Defaults to ten seconds.
    """

    metadata = None
    """
    An instance of :class:`cassandra.metadata.Metadata`.
    """

    connection_class = AsyncoreConnection
    """
    This determines what event loop system will be used for managing
    I/O with Cassandra.  These are the current options:

    * :class:`cassandra.io.asyncorereactor.AsyncoreConnection`
    * :class:`cassandra.io.libevreactor.LibevConnection`

    By default, ``AsyncoreConnection`` will be used, which uses
    the ``asyncore`` module in the Python standard library.  The
    performance is slightly worse than with ``libev``, but it is
    supported on a wider range of systems.
    """

    sessions = None
    control_connection = None
    scheduler = None
    executor = None
    _is_shutdown = False
    _is_setup = False
    _prepared_statements = None

    def __init__(self,
                 contact_points=("127.0.0.1",),
                 port=9042,
                 compression=True,
                 auth_provider=None,
                 load_balancing_policy=None,
                 reconnection_policy=None,
                 default_retry_policy=None,
                 conviction_policy_factory=None,
                 metrics_enabled=False,
                 connection_class=None,
                 sockopts=None,
                 executor_threads=2,
                 max_schema_agreement_wait=10):
        """
        Any of the mutable Cluster attributes may be set as keyword arguments
        to the constructor.
        """

        self.contact_points = contact_points
        self.port = port
        self.compression = compression

        if auth_provider is not None:
            if not callable(auth_provider):
                raise ValueError("auth_provider must be callable")
            self.auth_provider = auth_provider

        if load_balancing_policy is not None:
            self.load_balancing_policy = load_balancing_policy

        if reconnection_policy is not None:
            self.reconnection_policy = reconnection_policy

        if default_retry_policy is not None:
            self.default_retry_policy = default_retry_policy

        if conviction_policy_factory is not None:
            if not callable(conviction_policy_factory):
                raise ValueError("conviction_policy_factory must be callable")
            self.conviction_policy_factory = conviction_policy_factory

        if connection_class is not None:
            self.connection_class = connection_class

        self.metrics_enabled = metrics_enabled
        self.sockopts = sockopts
        self.max_schema_agreement_wait = max_schema_agreement_wait

        # let Session objects be GC'ed (and shutdown) when the user no longer
        # holds a reference. Normally the cycle detector would handle this,
        # but implementing __del__ prevents that.
        self.sessions = weakref.WeakSet()
        self.metadata = Metadata(self)
        self.control_connection = None
        self._prepared_statements = {}

        self._min_requests_per_connection = {
            HostDistance.LOCAL: DEFAULT_MIN_REQUESTS,
            HostDistance.REMOTE: DEFAULT_MIN_REQUESTS
        }

        self._max_requests_per_connection = {
            HostDistance.LOCAL: DEFAULT_MAX_REQUESTS,
            HostDistance.REMOTE: DEFAULT_MAX_REQUESTS
        }

        self._core_connections_per_host = {
            HostDistance.LOCAL: DEFAULT_MIN_CONNECTIONS_PER_LOCAL_HOST,
            HostDistance.REMOTE: DEFAULT_MIN_CONNECTIONS_PER_REMOTE_HOST
        }

        self._max_connections_per_host = {
            HostDistance.LOCAL: DEFAULT_MAX_CONNECTIONS_PER_LOCAL_HOST,
            HostDistance.REMOTE: DEFAULT_MAX_CONNECTIONS_PER_REMOTE_HOST
        }

        self.executor = ThreadPoolExecutor(max_workers=executor_threads)
        self.scheduler = _Scheduler(self.executor)

        self._lock = RLock()

        self.control_connection = ControlConnection(self)
        for address in contact_points:
            self.add_host(address, signal=True)

    def get_min_requests_per_connection(self, host_distance):
        return self._min_requests_per_connection[host_distance]

    def set_min_requests_per_connection(self, host_distance, min_requests):
        self._min_requests_per_connection[host_distance] = min_requests

    def get_max_requests_per_connection(self, host_distance):
        return self._max_requests_per_connection[host_distance]

    def set_max_requests_per_connection(self, host_distance, max_requests):
        self._max_requests_per_connection[host_distance] = max_requests

    def get_core_connections_per_host(self, host_distance):
        return self._core_connections_per_host[host_distance]

    def set_core_connections_per_host(self, host_distance, core_connections):
        old = self._core_connections_per_host[host_distance]
        self._core_connections_per_host[host_distance] = core_connections
        if old < core_connections:
            self.ensure_core_connections()

    def get_max_connections_per_host(self, host_distance):
        return self._max_connections_per_host[host_distance]

    def set_max_connections_per_host(self, host_distance, max_connections):
        self._max_connections_per_host[host_distance] = max_connections

    def connection_factory(self, address, *args, **kwargs):
        """
        Called to create a new connection with proper configuration.
        Intended for internal use only.
        """
        if self.auth_provider:
            kwargs['credentials'] = self.auth_provider(address)

        kwargs['compression'] = self.compression
        kwargs['sockopts'] = self.sockopts

        return self.connection_class.factory(address, *args, **kwargs)

    def _make_connection_factory(self, host, *args, **kwargs):
        if self.auth_provider:
            kwargs['credentials'] = self.auth_provider(host)

        kwargs['compression'] = self.compression
        kwargs['sockopts'] = self.sockopts

        return partial(self.connection_class.factory, host.address, *args, **kwargs)

    def connect(self, keyspace=None):
        """
        Creates and returns a new :class:`~.Session` object.  If `keyspace`
        is specified, that keyspace will be the default keyspace for
        operations on the ``Session``.
        """
        with self._lock:
            if self._is_shutdown:
                raise Exception("Cluster is already shut down")

            if not self._is_setup:
                self.load_balancing_policy.populate(
                    weakref.proxy(self), self.metadata.getAllHosts())
                self._is_setup = True

            if self.control_connection:
                try:
                    self.control_connection.connect()
                    log.debug("Control connection created")
                except:
                    log.exception("Control connection failed to connect, "
                                  "shutting down Cluster:")<commit_msg>Fix bad load balancing policy population<commit_after>"""
This module houses the main classes you will interact with,
:class:`.Cluster` and :class:`.Session`.
"""

from futures import ThreadPoolExecutor
import logging
import time
from threading import RLock, Thread, Event
import traceback
import Queue
import weakref
from functools import partial
from itertools import groupby

from cassandra import ConsistencyLevel, AuthenticationFailed
from cassandra.connection import ConnectionException
from cassandra.decoder import (QueryMessage, ResultMessage,
                               ErrorMessage, ReadTimeoutErrorMessage,
                               WriteTimeoutErrorMessage,
                               UnavailableErrorMessage,
                               OverloadedErrorMessage,
                               PrepareMessage, ExecuteMessage,
                               PreparedQueryNotFound,
                               IsBootstrappingErrorMessage, named_tuple_factory,
                               dict_factory)
from cassandra.io.asyncorereactor import AsyncoreConnection
from cassandra.metadata import Metadata
from cassandra.policies import (RoundRobinPolicy, SimpleConvictionPolicy,
                                ExponentialReconnectionPolicy, HostDistance,
                                RetryPolicy)
from cassandra.query import SimpleStatement, PreparedStatement, BoundStatement, bind_params
from cassandra.pool import (_ReconnectionHandler, _HostReconnectionHandler,
                            HostConnectionPool)


log = logging.getLogger(__name__)


DEFAULT_MIN_REQUESTS = 5
DEFAULT_MAX_REQUESTS = 100

DEFAULT_MIN_CONNECTIONS_PER_LOCAL_HOST = 2
DEFAULT_MAX_CONNECTIONS_PER_LOCAL_HOST = 8

DEFAULT_MIN_CONNECTIONS_PER_REMOTE_HOST = 1
DEFAULT_MAX_CONNECTIONS_PER_REMOTE_HOST = 2


class NoHostAvailable(Exception):
    """
    Raised when an operation is attempted but all connections are
    busy, defunct, closed, or resulted in errors when used.
    """

    errors = None
    """
    A map of the form ``{ip: exception}`` which details the particular
    Exception that was caught for each host the operation was attempted
    against.
    """

    def __init__(self, message, errors):
        Exception.__init__(self, message, errors)
        self.errors = errors


class Cluster(object):
    """
    The main class to use when interacting with a Cassandra cluster.
    Typically, one instance of this class will be created for each
    separate Cassandra cluster that your application interacts with.

    Example usage::

        >>> from cassandra.cluster import Cluster
        >>> cluster = Cluster(['192.168.1.1', '192.168.1.2'])
        >>> session = cluster.connect()
        >>> session.execute("CREATE KEYSPACE ...")
        >>> ...
        >>> cluster.shutdown()

    """

    port = 9042
    """
    The server-side port to open connections to. Defaults to 9042.
    """

    compression = True
    """
    Whether or not compression should be enabled when possible. Defaults to
    :const:`True` and attempts to use snappy compression.
    """

    auth_provider = None
    """
    An optional function that accepts one argument, the IP address of a node,
    and returns a dict of credentials for that node.
    """

    load_balancing_policy = RoundRobinPolicy()
    """
    An instance of :class:`.policies.LoadBalancingPolicy` or
    one of its subclasses.  Defaults to :class:`~.RoundRobinPolicy`.
    """

    reconnection_policy = ExponentialReconnectionPolicy(1.0, 600.0)
    """
    An instance of :class:`.policies.ReconnectionPolicy`. Defaults to an instance
    of :class:`.ExponentialReconnectionPolicy` with a base delay of one second and
    a max delay of ten minutes.
    """

    default_retry_policy = RetryPolicy()
    """
    A default :class:`.policies.RetryPolicy` instance to use for all
    :class:`.Query` objects which do not have a :attr:`~.Query.retry_policy`
    explicitly set.
    """

    conviction_policy_factory = SimpleConvictionPolicy
    """
    A factory function which creates instances of
    :class:`.policies.ConvictionPolicy`.  Defaults to
    :class:`.policies.SimpleConvictionPolicy`.
    """

    metrics_enabled = False
    """
    Whether or not metric collection is enabled.
    """

    sockopts = None
    """
    An optional list of tuples which will be used as arguments to
    ``socket.setsockopt()`` for all created sockets.
    """

    max_schema_agreement_wait = 10
    """
    The maximum duration (in seconds) that the driver will wait for schema
    agreement across the cluster. Defaults to ten seconds.
    """

    metadata = None
    """
    An instance of :class:`cassandra.metadata.Metadata`.
    """

    connection_class = AsyncoreConnection
    """
    This determines what event loop system will be used for managing
    I/O with Cassandra.  These are the current options:

    * :class:`cassandra.io.asyncorereactor.AsyncoreConnection`
    * :class:`cassandra.io.libevreactor.LibevConnection`

    By default, ``AsyncoreConnection`` will be used, which uses
    the ``asyncore`` module in the Python standard library.  The
    performance is slightly worse than with ``libev``, but it is
    supported on a wider range of systems.
    """

    sessions = None
    control_connection = None
    scheduler = None
    executor = None
    _is_shutdown = False
    _is_setup = False
    _prepared_statements = None

    def __init__(self,
                 contact_points=("127.0.0.1",),
                 port=9042,
                 compression=True,
                 auth_provider=None,
                 load_balancing_policy=None,
                 reconnection_policy=None,
                 default_retry_policy=None,
                 conviction_policy_factory=None,
                 metrics_enabled=False,
                 connection_class=None,
                 sockopts=None,
                 executor_threads=2,
                 max_schema_agreement_wait=10):
        """
        Any of the mutable Cluster attributes may be set as keyword arguments
        to the constructor.
        """

        self.contact_points = contact_points
        self.port = port
        self.compression = compression

        if auth_provider is not None:
            if not callable(auth_provider):
                raise ValueError("auth_provider must be callable")
            self.auth_provider = auth_provider

        if load_balancing_policy is not None:
            self.load_balancing_policy = load_balancing_policy

        if reconnection_policy is not None:
            self.reconnection_policy = reconnection_policy

        if default_retry_policy is not None:
            self.default_retry_policy = default_retry_policy

        if conviction_policy_factory is not None:
            if not callable(conviction_policy_factory):
                raise ValueError("conviction_policy_factory must be callable")
            self.conviction_policy_factory = conviction_policy_factory

        if connection_class is not None:
            self.connection_class = connection_class

        self.metrics_enabled = metrics_enabled
        self.sockopts = sockopts
        self.max_schema_agreement_wait = max_schema_agreement_wait

        # let Session objects be GC'ed (and shutdown) when the user no longer
        # holds a reference. Normally the cycle detector would handle this,
        # but implementing __del__ prevents that.
        self.sessions = weakref.WeakSet()
        self.metadata = Metadata(self)
        self.control_connection = None
        self._prepared_statements = {}

        self._min_requests_per_connection = {
            HostDistance.LOCAL: DEFAULT_MIN_REQUESTS,
            HostDistance.REMOTE: DEFAULT_MIN_REQUESTS
        }

        self._max_requests_per_connection = {
            HostDistance.LOCAL: DEFAULT_MAX_REQUESTS,
            HostDistance.REMOTE: DEFAULT_MAX_REQUESTS
        }

        self._core_connections_per_host = {
            HostDistance.LOCAL: DEFAULT_MIN_CONNECTIONS_PER_LOCAL_HOST,
            HostDistance.REMOTE: DEFAULT_MIN_CONNECTIONS_PER_REMOTE_HOST
        }

        self._max_connections_per_host = {
            HostDistance.LOCAL: DEFAULT_MAX_CONNECTIONS_PER_LOCAL_HOST,
            HostDistance.REMOTE: DEFAULT_MAX_CONNECTIONS_PER_REMOTE_HOST
        }

        self.executor = ThreadPoolExecutor(max_workers=executor_threads)
        self.scheduler = _Scheduler(self.executor)

        self._lock = RLock()

        self.control_connection = ControlConnection(self)
        for address in contact_points:
            self.add_host(address, signal=True)

    def get_min_requests_per_connection(self, host_distance):
        return self._min_requests_per_connection[host_distance]

    def set_min_requests_per_connection(self, host_distance, min_requests):
        self._min_requests_per_connection[host_distance] = min_requests

    def get_max_requests_per_connection(self, host_distance):
        return self._max_requests_per_connection[host_distance]

    def set_max_requests_per_connection(self, host_distance, max_requests):
        self._max_requests_per_connection[host_distance] = max_requests

    def get_core_connections_per_host(self, host_distance):
        return self._core_connections_per_host[host_distance]

    def set_core_connections_per_host(self, host_distance, core_connections):
        old = self._core_connections_per_host[host_distance]
        self._core_connections_per_host[host_distance] = core_connections
        if old < core_connections:
            self.ensure_core_connections()

    def get_max_connections_per_host(self, host_distance):
        return self._max_connections_per_host[host_distance]

    def set_max_connections_per_host(self, host_distance, max_connections):
        self._max_connections_per_host[host_distance] = max_connections

    def connection_factory(self, address, *args, **kwargs):
        """
        Called to create a new connection with proper configuration.
        Intended for internal use only.
        """
        if self.auth_provider:
            kwargs['credentials'] = self.auth_provider(address)

        kwargs['compression'] = self.compression
        kwargs['sockopts'] = self.sockopts

        return self.connection_class.factory(address, *args, **kwargs)

    def _make_connection_factory(self, host, *args, **kwargs):
        if self.auth_provider:
            kwargs['credentials'] = self.auth_provider(host)

        kwargs['compression'] = self.compression
        kwargs['sockopts'] = self.sockopts

        return partial(self.connection_class.factory, host.address, *args, **kwargs)

    def connect(self, keyspace=None):
        """
        Creates and returns a new :class:`~.Session` object.  If `keyspace`
        is specified, that keyspace will be the default keyspace for
        operations on the ``Session``.
        """
        with self._lock:
            if self._is_shutdown:
                raise Exception("Cluster is already shut down")

            if not self._is_setup:
                self.load_balancing_policy.populate(
                    weakref.proxy(self), self.metadata.all_hosts())
                self._is_setup = True

            if self.control_connection:
                try:
                    self.control_connection.connect()
                    log.debug("Control connection created")
                except:
                    log.exception("Control connection failed to connect, "
                                  "shutting down Cluster:")
{'fraction_non_alphanumeric': 0.07021949002121304, 'fraction_numerical': 0.00458894324429629, 'mean_word_length': 2.6014967259120674, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 6, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3996079', 'n_tokens_mistral': 6195, 'n_tokens_neox': 5880, 'n_words': 1594}