Source

python-hs-benchmark / handlersocket.py

Full commit
from pyhs import Manager
from MySQLdb import connect

import cPickle as pickle
import base64

DB_NAME = 'simple_test'
DB_USER = 'root'
DB_PASSWORD = ''
TABLE_NAME = 'test_table'

class Client(object):
    def __init__(self, dummyparam, engine='MYISAM'):
        connection = connect(
            host = '',
            user = DB_USER,
            passwd = DB_PASSWORD,
            db = DB_NAME,
            use_unicode = True
        )
                
        cursor = connection.cursor()
        
        # Stupid test to make test works in multi-threaded environment
        cursor.execute("SELECT count(ENGINE) FROM information_schema.TABLES "
            "WHERE lower(ENGINE) != '%s' AND TABLE_TYPE='BASE TABLE' "
            "AND TABLE_SCHEMA='%s'" % (engine.lower(), DB_NAME))
        
        # if test table already exists and it's engine differs from what we want
        # let's recreate it
        if cursor.fetchone()[0] > 0:
            cursor.execute('DROP TABLE %s' % TABLE_NAME)

            cursor.execute('CREATE TABLE `%s` (\
                `key` varchar(255) NOT NULL PRIMARY KEY,\
                `value` longtext NOT NULL\
            ) ENGINE = %s' % (TABLE_NAME, engine,))
        
        connection.close()
        
        self._index_fields = ['key', 'value',]
        self.manager = Manager( [('inet', 'localhost', 9998)],
                                [('inet', 'localhost', 9999)])
    
    def set(self, key, value):
        value = base64.encodestring(pickle.dumps(value, 2)).strip()
        
        if not self.has_key(key):
            self.manager.insert(
                    DB_NAME,
                    TABLE_NAME,
                    [   ('key', key), 
                        ('value', value)]
            )
        else:
            self.manager.update(
                DB_NAME,
                TABLE_NAME,
                '=', 
                self._index_fields,
                [key],
                [   key, 
                    value]
            )
            
    def has_key(self, key):
        raw = self.manager.get(DB_NAME, TABLE_NAME, 
                ['key'], key)
        
        return bool(len(raw))
    
    
    def get(self, key):
        raw = self.manager.get(DB_NAME, TABLE_NAME, 
                self._index_fields, key)
        
        if len(raw) == 0:
            return None
        else:
            return pickle.loads(base64.decodestring(raw[1][1]))

    def get_multi(self, keys):
        return dict([(key, self.get(key)) for key in keys])

    def delete(self, key):
        self.manager.delete(
            DB_NAME,
            TABLE_NAME,
            '=', 
            self._index_fields,
            [key]
        )