small improvements

Gravelfile

@@ -1,5 +1,5 @@
 name: gravel-common
-requires-apt: python2.7 python-tdb python-yaml libpcre3 python-bson-ext python-passfd
+requires-apt: python2.7 python-tdb python-yaml libpcre3 python-bson python-passfd
 symlinks:
 - [dbtool.py, /usr/local/bin/graveldbtool]
 - [rpctool.py, /usr/local/bin/gravelrpctool]

graveldb.py

@@ -43,6 +43,8 @@ def __exit__(self, type, value, tb):
 def _mkkey(k):
     if isinstance(k, (str, int)):
         return str(k)
+    elif isinstance(k, unicode):
+        return k.encode('utf8')
     else:
         raise TypeError('expected str or int')
 
@@ -58,20 +60,21 @@ def Table(name, path):
     return type(name, (_Table,), {'table': TDB_BSON_Shelf(path + '/' + name)})
 
 class _Table(object):
+    __slots__ = ('data', 'name', 'exists')
     autocreate = True
 
     def __init__(self, name, autocreate=None):
         self.name = name
+        if autocreate is None:
+            autocreate = self.autocreate
         try:
             self.data = self.table[name]
             self.exists = True
-        except:
+        except KeyError:
+            if not autocreate:
+                raise
             self.data = Object()
             self.exists = False
-        if autocreate is None:
-            autocreate = self.autocreate
-        if not self.exists and not autocreate:
-            raise KeyError(name)
         self._setup()
 
     def _setup(self):
@@ -101,6 +104,9 @@ def __enter__(self):
     def __exit__(self, type, value, tb):
         return self.table.__exit__(type, value, tb)
 
+    def __repr__(self):
+        return '<%s %s: %r>' % (self.__class__.__name__, self.name, self.data.__dict__)
+
     default = {}
 
 class Object(object):

gravelrpc.py

@@ -1,12 +1,12 @@
 import SocketServer
 import socket
 import functools
-import bson
 import os
 import traceback
 import passfd
 import struct
 
+import bson as _bson
 from bson.binary import Binary
 
 PATH = '/gravel/run/%s.sock'
@@ -82,7 +82,7 @@ def _rpc_write_bson(sock, doc):
             raise TypeError('fds need to be instances of FD (not %r)' % fd)
         passfd.sendfd(sock, fd.fileno(), 'whatever')
 
-    sock.sendall(bson.BSON.encode(doc))
+    sock.sendall(_bson.BSON.encode(doc))
     sock.shutdown(socket.SHUT_WR)
 
 def _rpc_read_bson(sock, allow_fd_passing=False):
@@ -94,17 +94,27 @@ def _rpc_read_bson(sock, allow_fd_passing=False):
         raise IOError('client tried to pass fds')
 
     raw = ''.join(iter(lambda: sock.recv(4096), ''))
-    result = bson.BSON(raw).decode()
+    result = _bson.BSON(raw).decode()
     if fd_count != 0:
         result['fds'] = fds
     elif 'fds' in result:
         del result['fds']
     return result
 
-def bson_load(f):
-    length_data = f.read(4)
-    length, = struct.unpack('<I', length_data)
-    return bson.BSON(length_data + f.read(length - 4)).decode()
+class bson:
+    ''' pickle/marshal/json compatiblity module for BSON '''
+    def load(self, f):
+        length_data = f.read(4)
+        length, = struct.unpack('<I', length_data)
+        return _bson.BSON(length_data + f.read(length - 4)).decode()
 
-def bson_dump(obj, f):
-    f.write(bson.BSON.encode(obj))
+    def dump(self, obj, f):
+        f.write(_bson.BSON.encode(obj))
+
+    def loads(self, s):
+        return _bson.BSON(s).decode()
+
+    def dumps(self, obj):
+        return _bson.BSON.encode(obj)
+
+bson = bson()

ssh_utils.py

@@ -57,16 +57,24 @@ def write_authorized_keys(keys, user=''):
     os.rename(out.name, path)
 
 def call(address, *args, **kwargs):
-    key = kwargs.get('key')
+    def _get_kwargs(key=None, decode=None, check_output=True):
+        return key, decode, check_output
+
+    key, decode, check_output = _get_kwargs(**kwargs)
     if ':' in address:
         hostname, port = address.split(':')
     else:
         hostname = address
         port = None
-    opts = ['-o', 'StrictHostKeyChecking=no']
+    opts = ['-o', 'StrictHostKeyChecking=no', '-o', 'PasswordAuthentication=no']
     if key:
         opts += ['-i', key]
     if port:
         opts += ['-p', port]
-    return subprocess.check_output([
+    func = subprocess.check_output if check_output else subprocess.check_call
+    result = func([
         'ssh', ] + opts + [hostname, '--'] + list(args))
+    if decode:
+        return decode.loads(result)
+    else:
+        return result