<commit_before># -*- test-case-name: vumi.persist.tests.test_txriak_manager -*-

"""A manager implementation on top of txriak."""

from riakasaurus.riak import RiakClient, RiakObject, RiakMapReduce
from twisted.internet.defer import (
    inlineCallbacks, gatherResults, maybeDeferred, succeed)

from vumi.persist.model import Manager


class TxRiakManager(Manager):
    """A persistence manager for txriak."""

    call_decorator = staticmethod(inlineCallbacks)

    @classmethod
    def from_config(cls, config):
        config = config.copy()
        bucket_prefix = config.pop('bucket_prefix')
        client = RiakClient(**config)
        return cls(client, bucket_prefix)

    def riak_object(self, cls, key, result=None):
        bucket = self.bucket_for_cls(cls)
        riak_object = RiakObject(self.client, bucket, key)
        if result:
            metadata = result['metadata']
            indexes = metadata['index']
            if hasattr(indexes, 'items'):
                # TODO: I think this is a Riak bug. In some cases
                #       (maybe when there are no indexes?) the index
                #       comes back as a list, in others (maybe when
                #       there are indexes?) it comes back as a dict.
                indexes = indexes.items()

            content_type = metadata['content-type'].encode('utf-8')
            indexes = [
                (k.encode('utf-8'), v.encode('utf-8')) for k, v in indexes]
            data = result['data'].encode('utf-8')

            riak_object.set_content_type(content_type)<commit_msg>encode Indexes better when given tuples<commit_after># -*- test-case-name: vumi.persist.tests.test_txriak_manager -*-

"""A manager implementation on top of txriak."""

from riakasaurus.riak import RiakClient, RiakObject, RiakMapReduce
from twisted.internet.defer import (
    inlineCallbacks, gatherResults, maybeDeferred, succeed)

from vumi.persist.model import Manager


class TxRiakManager(Manager):
    """A persistence manager for txriak."""

    call_decorator = staticmethod(inlineCallbacks)

    @classmethod
    def from_config(cls, config):
        config = config.copy()
        bucket_prefix = config.pop('bucket_prefix')
        client = RiakClient(**config)
        return cls(client, bucket_prefix)

    def _encode_indexes(self, iterable, encoding='utf-8'):
        """
        From Basho's docs:

            When using the HTTP interface, multi-valued indexes are specified
            by separating the values with a comma (,). For that reason,
            your application should avoid using a comma as part of an
            index value.

        The index values we get can either be a single string value or can
        be a tuple of multiple values that need to be set. If we get a tuple
        then convert it to a comma separated string.
        """
        encoded = []
        for key, value in iterable:
            if isinstance(value, (list, tuple)):
                value = ", ".join([v.encode(encoding) for v in value])
                key = key.encode(encoding)
            else:
                value = value.encode(encoding)
                key = key.encode(encoding)
                encoded.append((key, value))
        return encoded

    def riak_object(self, cls, key, result=None):
        bucket = self.bucket_for_cls(cls)
        riak_object = RiakObject(self.client, bucket, key)
        if result:
            metadata = result['metadata']
            indexes = metadata['index']
            if hasattr(indexes, 'items'):
                # TODO: I think this is a Riak bug. In some cases
                #       (maybe when there are no indexes?) the index
                #       comes back as a list, in others (maybe when
                #       there are indexes?) it comes back as a dict.
                indexes = indexes.items()

            content_type = metadata['content-type'].encode('utf-8')
            indexes = self._encode_indexes(indexes, 'utf-8')
            data = result['data'].encode('utf-8')

            riak_object.set_content_type(content_type)
{'fraction_non_alphanumeric': 0.09652981427174975, 'fraction_numerical': 0.0019550342130987292, 'mean_word_length': 2.245836637589215, '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': '1339545', 'n_tokens_mistral': 1157, 'n_tokens_neox': 1095, 'n_words': 345}