#!/usr/bin/env python"""Copyright 2009 Jay Baird <jay@mochimedia.com>This file is provided to you under the Apache License,Version 2.0 (the "License"); you may not use this fileexcept in compliance with the License. You may obtaina copy of the License at http://www.apache.org/licenses/LICENSE-2.0Unless required by applicable law or agreed to in writing,software distributed under the License is distributed on an"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANYKIND, either express or implied. See the License for thespecific language governing permissions and limitationsunder the License."""importurllibfromcStringIOimportStringIOtry:importpycurlHAS_PYCURL=TrueexceptImportError:importhttplibHAS_PYCURL=Falsetry:importjsonexceptImportError:importsimplejsonasjsondefexpect(status):"""Wraps a function in a function that guarantees a return code(s) and if certain conditions are met either return None or a Dictionary from a JSON string. """ifnotisinstance(status,(list,tuple)):status=[status]defwrapper_func(f):defwrapped_func(*args,**kwargs):code,resp=f(*args,**kwargs)ifHAS_PYCURL:resp.reset()ifcodenotinstatus:raiseJiakException(code,resp.read())ifcodein[404,204]:returnNoneelse:returnjson.load(resp)returnwrapped_funcreturnwrapper_funcdefbuild_headers(h,disable_continue=True):headers=['Expect:']ifdisable_continueelse[]fork,vinh.iteritems():headers.append('%s: %s'%(k,v))returnheadersclassJiakException(Exception):passclassJiak(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 something more Pythonic and I wanted to use pycURL. Example Usage: >>> client = Jiak('127.0.0.1', 8098, 'jiak') >>> [client.delete('jiak_example', key) for key in ['doctestkey', 'jroot', 'jleaf1', 'jleaf2', 'jleaf3']] [None, None, None, None, None] >>> obj = client.store('jiak_example', 'doctestkey', {'foo':2}) >>> client.fetch('jiak_example', 'doctestkey').get('object', {}).get('foo') == 2 True """def__init__(self,host,port,prefix="jiak"):self.host=hostself.port=portself.prefix=prefixifHAS_PYCURL:self._request=self._pycurl_requestelse:self._request=self._httplib_requestdef_build_path(self,bucket,key=''):return'http://%s:%d/%s/%s/%s'%(self.host,self.port,self.prefix,urllib.quote_plus(bucket),urllib.quote_plus(key))def_httplib_request(self,method,uri,body="",headers={}):client=httplib.HTTPConnection(self.host,self.port)client.request(method,uri,body,headers)response=client.getresponse()returnresponse.status,response.getheaders(),responsedef_pycurl_request(self,method,uri,body="",headers={}):resp_headers=StringIO()response=StringIO()client=pycurl.Curl()ifmethodin("PUT","POST"):ifmethod=="POST":client.setopt(pycurl.POST,1)else:client.setopt(pycurl.CUSTOMREQUEST,method)client.setopt(pycurl.POSTFIELDS,body)elifmethodin("DELETE",):client.setopt(pycurl.CUSTOMREQUEST,method)client.setopt(pycurl.URL,uri)client.setopt(pycurl.HTTPHEADER,build_headers(headers))client.setopt(pycurl.WRITEFUNCTION,response.write)client.setopt(pycurl.HEADERFUNCTION,resp_headers.write)client.perform()code=client.getinfo(pycurl.HTTP_CODE)returncode,resp_headers,response@expect(204)defset_bucket_schema(self,bucket,allowed_fields,required_fields=[],write_mask=None,read_mask=None):write_mask=allowed_fieldsifwrite_maskisNoneelsewrite_maskread_mask=allowed_fieldsifread_maskisNoneelseread_maskbody=json.dumps(dict(schema=dict(allowed_fields=allowed_fields,required_fields=required_fields,write_mask=write_mask,read_mask=read_mask)))code,_,resp=self._request("PUT",self._build_path(bucket),body,{'Content-Type':"application/json"})returncode,resp@expect(200)deflist_bucket(self,bucket):code,_,resp=self._request("GET",self._build_path(bucket))returncode,resp@expect(200)defstore(self,bucket,key,obj,links=[],w=2,dw=2):obj=dict(bucket=bucket,key=key,object=obj,links=links)code,_,resp=self._request("PUT",'%s?%s'%(self._build_path(bucket,key),urllib.urlencode(dict(returnbody='true',w=w,dw=dw))),json.dumps(obj),{'Content-Type':'application/json'})returncode,resp@expect([200,404])deffetch(self,bucket,key,r=2):code,headers,resp=self._request("GET",'%s?r=%d'%(self._build_path(bucket,key),r))returncode,resp@expect([204,404])defdelete(self,bucket,key,dw=2):code,_,resp=self._request("DELETE",'%s?dw=%d'%(self._build_path(bucket,key),dw))returncode,resp@expect([200,404])defwalk(self,bucket,key,spec):"""spec should be a list of tuples, each of the form: (bucket, tag, acc) where bucket is a string name of a bucket, or "_" to match any bucket tag is a string tag name, or "_" to match any link tag acc is either the string "1" or "0" if the walk succeeds, this will return a list, where each element is a list of JiakObjects corresponding to a spec element that had acc == "1" """defbuild_spec(spec):return"/".join(['%s,%s,%s'%tuple(map(urllib.quote_plus,[b,t,a]))forb,t,ainspec])code,_,resp=self._request("GET",'%s/%s'%(self._build_path(bucket,key),build_spec(spec)))returncode,respif__name__=='__main__':importdoctestdoctest.testmod()