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

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

from txriak.riak import RiakClient, RiakObject, RiakMapReduce, RiakLink
from twisted.internet.defer import inlineCallbacks, DeferredList

from vumi.persist.model import Manager


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

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

    def riak_object(self, modelobj):
        bucket_name = self.bucket_prefix + modelobj.bucket
        bucket = self.client.bucket(bucket_name)
        riak_object = RiakObject(self.client, bucket, modelobj.key)
        riak_object.set_data({})
        riak_object.set_content_type("application/json")
        return riak_object

    def store(self, modelobj):
        d = modelobj._riak_object.store()
        d.addCallback(lambda result: modelobj)
        return d

    def load(self, modelobj):
        d = modelobj._riak_object.reload()
        d.addCallback(lambda result: modelobj
                      if result.get_data() is not None else None)
        return d

    def load_list(self, modelobjs):
        deferreds = []
        for modelobj in modelobjs:
            deferreds.append(self.load(modelobj))
        d = DeferredList(deferreds)

        def strip_results(mapped_results):
            return [t[1] for t in mapped_results]

        d.addCallback(strip_results)
        return d

    def riak_map_reduce(self):
        return RiakMapReduce(self.client)

    def run_map_reduce(self, mapreduce, mapper_func):
        mapreduce_done = mapreduce.run()

        def map_results(raw_results):
            deferreds = []
            for row in raw_results:
                link = RiakLink(row[0], row[1])
                deferreds.append(mapper_func(self, link))
            return DeferredList(deferreds)

        def strip_results(mapped_results):
            return [t[1] for t in mapped_results]

        mapreduce_done.addCallback(map_results)
        mapreduce_done.addCallback(strip_results)

        return mapreduce_done

    @inlineCallbacks
    def purge_all(self):
        buckets = yield self.client.list_buckets()
        for bucket_name in buckets:
            if bucket_name.startswith(self.bucket_prefix):
                bucket = self.client.bucket(bucket_name)
                yield bucket.purge_keys()
<commit_msg>Replace DeferredLists with gatherResults (which does exactly what we want, unlike DeferredList which requires a bunch of extra code on our side.<commit_after># -*- test-case-name: vumi.persist.tests.test_txriak_manager -*-

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

from txriak.riak import RiakClient, RiakObject, RiakMapReduce, RiakLink
from twisted.internet.defer import inlineCallbacks, gatherResults

from vumi.persist.model import Manager


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

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

    def riak_object(self, modelobj):
        bucket_name = self.bucket_prefix + modelobj.bucket
        bucket = self.client.bucket(bucket_name)
        riak_object = RiakObject(self.client, bucket, modelobj.key)
        riak_object.set_data({})
        riak_object.set_content_type("application/json")
        return riak_object

    def store(self, modelobj):
        d = modelobj._riak_object.store()
        d.addCallback(lambda result: modelobj)
        return d

    def load(self, modelobj):
        d = modelobj._riak_object.reload()
        d.addCallback(lambda result: modelobj
                      if result.get_data() is not None else None)
        return d

    def load_list(self, modelobjs):
        deferreds = []
        for modelobj in modelobjs:
            deferreds.append(self.load(modelobj))
        return gatherResults(deferreds)

    def riak_map_reduce(self):
        return RiakMapReduce(self.client)

    def run_map_reduce(self, mapreduce, mapper_func):
        mapreduce_done = mapreduce.run()

        def map_results(raw_results):
            deferreds = []
            for row in raw_results:
                link = RiakLink(row[0], row[1])
                deferreds.append(mapper_func(self, link))
            return gatherResults(deferreds)

        mapreduce_done.addCallback(map_results)
        return mapreduce_done

    @inlineCallbacks
    def purge_all(self):
        buckets = yield self.client.list_buckets()
        for bucket_name in buckets:
            if bucket_name.startswith(self.bucket_prefix):
                bucket = self.client.bucket(bucket_name)
                yield bucket.purge_keys()

{'fraction_non_alphanumeric': 0.08727719729563614, 'fraction_numerical': 0.001229256299938537, 'mean_word_length': 2.659670164917541, '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': '1054514', 'n_tokens_mistral': 1485, 'n_tokens_neox': 1423, 'n_words': 323}