Commits

Anonymous committed 1f6e824

* Add X-Riak-ClientId header to requests using similar logic to jiak.rb and jiak.js
* Removed dicts as default args due to probable undesirable behavior, see http://magicrebirth.wordpress.com/2009/05/18/python-pitfall-passing-mutable-objects-as-default-args/

  • Participants
  • Parent commits 57d4a53

Comments (0)

Files changed (1)

client_lib/jiak.py

 
 class JiakException(Exception): pass
 
+MAX_CLIENT_ID = 4294967296
+CLIENT_ID_HDR = 'X-Riak-ClientId'
+
+# base64-encode a 32-bit integer
+def int_to_base64(n):
+    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
+    return "%c%c%c%c%c%c==" % (
+        alphabet[(n >> 26)],
+        alphabet[((n >> 20)&63)],
+        alphabet[((n >> 14)&63)],
+        alphabet[((n >> 8)&63)],
+        alphabet[((n >> 2)&63)],
+        alphabet[((n << 4)&63)])
+
 class Jiak(object):
     """A Python interface for the Riak (http://riak.basho.com/) key-value store.
        The Riak source does ship with a client library for python, but I wanted
        True
        
     """
-    def __init__(self, host, port, prefix="jiak"):
+    def __init__(self, host, port, prefix="jiak", client_id=None):
         self.host = host
         self.port = port
         self.prefix = prefix
+        self.headers = {}
+
+        if client_id is not None:
+            if 0 < client_id < MAX_CLIENT_ID:
+                self.headers[CLIENT_ID_HDR] = int_to_base64(client_id)
+        else:
+            self.headers[CLIENT_ID_HDR] = int_to_base64(random.randint(0, MAX_CLIENT_ID))
+
         if HAS_PYCURL:
             self._request = self._pycurl_request
         else:
                                             urllib.quote_plus(bucket),
                                             urllib.quote_plus(key))
     
-    def _httplib_request(self, method, uri, body="", headers={}):
+    def _httplib_request(self, method, uri, body="", hdrs=None):
+        headers = self.headers.copy()
+        if hdrs is not None:
+            headers.update(hdrs)
         client = httplib.HTTPConnection(self.host, self.port)
         client.request(method, uri, body, headers)
         response = client.getresponse()
         return response.status, response.getheaders(), response
         
-    def _pycurl_request(self, method, uri, body="", headers={}):
+    def _pycurl_request(self, method, uri, body="", hdrs=None):
+        headers = self.headers.copy()
+        if hdrs is not None:
+            headers.update(hdrs)
         resp_headers = StringIO()
         response = StringIO()
         client = pycurl.Curl()
 
 if __name__ == '__main__':
     import doctest
-    doctest.testmod()
+    doctest.testmod()