Merge pull request #106 from ultimatebuster/setlinks

Issue #83 fix: added set_links

riak/riak_object.py

@@ -212,6 +212,22 @@ def remove_index(self, field, value):
             self._metadata[MD_INDEX].remove(rie)
         return self
 
+    def set_indexes(self, indexes):
+        """
+        Sets indexes once and for all. Currenly supports an iterable of 2 item tuples,
+        (field, value)
+
+        :param indexes: iterable of 2 item tuples consisting the field and value.
+        :rtype: self
+        """
+        new_indexes = []
+        for field, value in indexes:
+            rie = RiakIndexEntry(field, value)
+            new_indexes.append(rie)
+        self._metadata[MD_INDEX] = new_indexes
+
+        return self
+
     def get_indexes(self, field = None):
         """
         Get a list of the index entries for this object. If a field is provided, returns a list
@@ -263,6 +279,33 @@ def set_content_type(self, content_type):
         self._metadata[MD_CTYPE] = content_type
         return self
 
+    def set_links(self, links, all_link=False):
+        """
+        Replaces all links to a RiakObject
+
+        :param links: An iterable of 2-item tuples, consisting of (RiakObject, tag). This could also be an iterable of
+            just a RiakObject, instead of the tuple, then a tag of None would be used. Lastly, it could also be an
+            iterable of RiakLink. They have tags built-in.
+        :param all_link: A boolean indicates if links is all RiakLink object
+            This speeds up the operation.
+        """
+        if all_link:
+            self._metadata[MD_LINKS] = links
+            return self
+
+        new_links = []
+        for item in links:
+            if isinstance(item, RiakLink):
+                link = item
+            elif isinstance(item, RiakObject):
+                link = RiakLink(item._bucket._name, item._key, None)
+            else:
+                link = RiakLink(item[0]._bucket._name, item[0]._key, item[1])
+            new_links.append(link)
+
+        self._metadata[MD_LINKS] = new_links
+        return self
+
     def add_link(self, obj, tag=None):
         """
         Add a link to a RiakObject.

riak/tests/test_all.py

@@ -451,6 +451,31 @@ def test_store_and_get_links(self):
             else:
                 self.assertEqual("unknown key", l.get_key())
 
+    def test_set_links(self):
+        # Create the object
+        bucket = self.client.bucket("bucket")
+        bucket.new("foo", 2).set_links([bucket.new("foo1"),
+            (bucket.new("foo2"), "tag"),
+            RiakLink("bucket", "foo2", "tag2")]).store()
+        obj = bucket.get("foo")
+        links = sorted(obj.get_links(), key=lambda x: x.get_key())
+        self.assertEqual(len(links), 3)
+        self.assertEqual(links[0].get_key(), "foo1")
+        self.assertEqual(links[1].get_key(), "foo2")
+        self.assertEqual(links[1].get_tag(), "tag")
+        self.assertEqual(links[2].get_key(), "foo2")
+        self.assertEqual(links[2].get_tag(), "tag2")
+
+    def test_set_links_all_links(self):
+        bucket = self.client.bucket("bucket")
+        foo1 = bucket.new("foo", 1)
+        foo2 = bucket.new("foo2", 2).store()
+        links = [RiakLink("bucket", "foo2")]
+        foo1.set_links(links, True)
+        links = foo1.get_links()
+        self.assertEqual(len(links), 1)
+        self.assertEqual(links[0].get_key(), "foo2")
+
     def test_link_walking(self):
         # Create the object...
         bucket = self.client.bucket("bucket")
@@ -630,6 +655,22 @@ def test_secondary_index_store(self):
         # Clean up...
         bucket.get('mykey1').delete()
 
+    @unittest.skipIf(SKIP_INDEXES, 'SKIP_INDEXES is defined')
+    def test_set_indexes(self):
+        if not self.is_2i_supported():
+            return True
+
+        bucket = self.client.bucket('indexbucket')
+        foo = bucket.new('foo', 1)
+        foo.set_indexes((('field1_bin', 'test'), ('field2_int', 1337))).store()
+        result = self.client.index('indexbucket', 'field2_int', 1337).run()
+        self.assertEqual(1, len(result))
+        self.assertEqual('foo', result[0].get_key())
+
+        result = self.client.index('indexbucket', 'field1_bin', 'test').run()
+        self.assertEqual(1, len(result))
+        self.assertEqual('foo', result[0].get_key())
+
     @unittest.skipIf(SKIP_INDEXES, 'SKIP_INDEXES is defined')
     def test_secondary_index_query(self):
         if not self.is_2i_supported():