Permalink
Fetching contributors…
Cannot retrieve contributors at this time
130 lines (104 sloc) 4.14 KB
#!/usr/bin/env python
"""
TAP protocol client library.
Copyright (c) 2010 Dustin Sallings <dustin@spy.net>
"""
import socket
import string
import random
import struct
import asyncore
import mc_bin_server
import mc_bin_client
from memcacheConstants import REQ_MAGIC_BYTE, RES_MAGIC_BYTE
from memcacheConstants import REQ_PKT_FMT, RES_PKT_FMT, MIN_RECV_PACKET
from memcacheConstants import SET_PKT_FMT, DEL_PKT_FMT, INCRDECR_RES_FMT
import memcacheConstants
class TapConnection(mc_bin_server.MemcachedBinaryChannel):
def __init__(self, server, port, callback, clientId=None, opts={}, user=None, pswd=None):
mc_bin_server.MemcachedBinaryChannel.__init__(self, None, None,
self._createTapCall(clientId,
opts))
self.server = server
self.port = port
self.callback = callback
self.identifier = (server, port)
self.user = user
self.pswd = pswd
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
self.connect((server, port))
def create_socket(self, family, type):
if not self.user:
mc_bin_server.MemcachedBinaryChannel.create_socket(self, family, type)
return
self.family_and_type = family, type
self.mc = mc_bin_client.MemcachedClient(self.server, self.port)
self.mc.sasl_auth_plain(self.user, self.pswd or "")
sock = self.mc.s
sock.setblocking(0)
self.set_socket(sock)
def _createTapCall(self, key=None, opts={}):
# Client identifier
if not key:
key = "".join(random.sample(string.letters, 16))
dtype=0
opaque=0
cas=0
extraHeader, val = self._encodeOpts(opts)
msg=struct.pack(REQ_PKT_FMT, REQ_MAGIC_BYTE,
memcacheConstants.CMD_TAP_CONNECT,
len(key), len(extraHeader), dtype, 0,
len(key) + len(extraHeader) + len(val),
opaque, cas)
return msg + extraHeader + key + val
def _encodeOpts(self, opts):
header = 0
val = []
for op in sorted(opts.keys()):
header |= op
if op in memcacheConstants.TAP_FLAG_TYPES:
val.append(struct.pack(memcacheConstants.TAP_FLAG_TYPES[op],
opts[op]))
elif op == memcacheConstants.TAP_FLAG_LIST_VBUCKETS:
val.append(self._encodeVBucketList(opts[op]))
else:
val.append(opts[op])
return struct.pack(">I", header), ''.join(val)
def _encodeVBucketList(self, vbl):
l = list(vbl) # in case it's a generator
vals = [struct.pack("!H", len(l))]
for v in vbl:
vals.append(struct.pack("!H", v))
return ''.join(vals)
def processCommand(self, cmd, klen, vb, extralen, cas, data):
extra = data[0:extralen]
if len(extra) == 8:
es_length, _flags, _ttl = \
struct.unpack(memcacheConstants.TAP_GENERAL_PKT_FMT, extra)
else:
es_length, _flags, _ttl, _flg, _exp = \
struct.unpack(memcacheConstants.TAP_MUTATION_PKT_FMT, extra)
key_start = extralen + es_length
key = data[key_start:(key_start+klen)]
val = data[(key_start+klen):]
return self.callback(self.identifier, cmd, extra, key, vb, val, cas)
def handle_connect(self):
pass
def handle_close(self):
self.close()
class TapClient(object):
def __init__(self, servers, callback, opts={}, user=None, pswd=None):
for t in servers:
tc = TapConnection(t.host, t.port, callback, t.id, opts, user, pswd)
class TapDescriptor(object):
port = 11211
id = None
def __init__(self, s):
self.host = s
if ':' in s:
self.host, self.port = s.split(':', 1)
self.port = int(self.port)
if '@' in self.host:
self.id, self.host = self.host.split('@', 1)
def __repr__(self):
return "<TapDescriptor %s@%s:%d>" % (self.id or "(anon)", self.host, self.port)