IPython/zmq/parallel/sqlitedb.py

"""A TaskRecord backend using sqlite3"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-----------------------------------------------------------------------------

import json
import os
import cPickle as pickle
from datetime import datetime

import sqlite3

from IPython.utils.traitlets import CUnicode, CStr, Instance, List
from .dictdb import BaseDB
from .util import ISO8601

#-----------------------------------------------------------------------------
# SQLite operators, adapters, and converters
#-----------------------------------------------------------------------------

operators = {
 '$lt' : lambda a,b: "%s < ?",
 '$gt' : ">",
 # null is handled weird with ==,!=
 '$eq' : "IS",
 '$ne' : "IS NOT",
 '$lte': "<=",
 '$gte': ">=",
 '$in' : ('IS', ' OR '),
 '$nin': ('IS NOT', ' AND '),
 # '$all': None,
 # '$mod': None,
 # '$exists' : None
}

def _adapt_datetime(dt):
    return dt.strftime(ISO8601)

def _convert_datetime(ds):
    if ds is None:
        return ds
    else:
        return datetime.strptime(ds, ISO8601)

def _adapt_dict(d):
    return json.dumps(d)

def _convert_dict(ds):
    if ds is None:
        return ds
    else:
        return json.loads(ds)

def _adapt_bufs(bufs):
    # this is *horrible*
    # copy buffers into single list and pickle it:
    if bufs and isinstance(bufs[0], (bytes, buffer)):
        return sqlite3.Binary(pickle.dumps(map(bytes, bufs),-1))
    elif bufs:
        return bufs
    else:
        return None

def _convert_bufs(bs):
    if bs is None:
        return []
    else:
        return pickle.loads(bytes(bs))

#-----------------------------------------------------------------------------
# SQLiteDB class
#-----------------------------------------------------------------------------

class SQLiteDB(BaseDB):
    """SQLite3 TaskRecord backend."""
    
    filename = CUnicode('tasks.db', config=True)
    location = CUnicode('', config=True)
    table = CUnicode("", config=True)
    
    _db = Instance('sqlite3.Connection')
    _keys = List(['msg_id' ,
            'header' ,
            'content',
            'buffers',
            'submitted',
            'client_uuid' ,
            'engine_uuid' ,
            'started',
            'completed',
            'resubmitted',
            'result_header' ,
            'result_content' ,
            'result_buffers' ,
            'queue' ,
            'pyin' ,
            'pyout',
            'pyerr',
            'stdout',
            'stderr',
        ])
    
    def __init__(self, **kwargs):
        super(SQLiteDB, self).__init__(**kwargs)
        if not self.table:
            # use session, and prefix _, since starting with # is illegal
            self.table = '_'+self.session.replace('-','_')
        if not self.location:
            if hasattr(self.config.Global, 'cluster_dir'):
                self.location = self.config.Global.cluster_dir
            else:
                self.location = '.'
        self._init_db()
    
    def _defaults(self):
        """create an empty record"""
        d = {}
        for key in self._keys:
            d[key] = None
        return d
    
    def _init_db(self):
        """Connect to the database and get new session number."""
        # register adapters
        sqlite3.register_adapter(datetime, _adapt_datetime)
        sqlite3.register_converter('datetime', _convert_datetime)
        sqlite3.register_adapter(dict, _adapt_dict)
        sqlite3.register_converter('dict', _convert_dict)
        sqlite3.register_adapter(list, _adapt_bufs)
        sqlite3.register_converter('bufs', _convert_bufs)
        # connect to the db
        dbfile = os.path.join(self.location, self.filename)
        self._db = sqlite3.connect(dbfile, detect_types=sqlite3.PARSE_DECLTYPES)
        
        self._db.execute("""CREATE TABLE IF NOT EXISTS %s 
                (msg_id text PRIMARY KEY,
                header dict text,
                content dict text,
                buffers bufs blob,
                submitted datetime text,
                client_uuid text,
                engine_uuid text,
                started datetime text,
                completed datetime text,
                resubmitted datetime text,
                result_header dict text,
                result_content dict text,
                result_buffers bufs blob,
                queue text,
                pyin text,
                pyout text,
                pyerr text,
                stdout text,
                stderr text)
                """%self.table)
        # self._db.execute("""CREATE TABLE IF NOT EXISTS %s_buffers 
        #         (msg_id text, result integer, buffer blob)
        #         """%self.table)
        self._db.commit()
    
    def _dict_to_list(self, d):
        """turn a mongodb-style record dict into a list."""
        
        return [ d[key] for key in self._keys ]
    
    def _list_to_dict(self, line):
        """Inverse of dict_to_list"""
        d = self._defaults()
        for key,value in zip(self._keys, line):
            d[key] = value
        
        return d
    
    def _render_expression(self, check):
        """Turn a mongodb-style search dict into an SQL query."""
        expressions = []
        args = []
        
        skeys = set(check.keys())
        skeys.difference_update(set(self._keys))
        skeys.difference_update(set(['buffers', 'result_buffers']))
        if skeys:
            raise KeyError("Illegal testing key(s): %s"%skeys)
        
        for name,sub_check in check.iteritems():
            if isinstance(sub_check, dict):
                for test,value in sub_check.iteritems():
                    try:
                        op = operators[test]
                    except KeyError:
                        raise KeyError("Unsupported operator: %r"%test)
                    if isinstance(op, tuple):
                        op, join = op
                    expr = "%s %s ?"%(name, op)
                    if isinstance(value, (tuple,list)):
                        expr = '( %s )'%( join.join([expr]*len(value)) )
                        args.extend(value)
                    else:
                        args.append(value)
                    expressions.append(expr)
            else:
                # it's an equality check
                expressions.append("%s IS ?"%name)
                args.append(sub_check)
        
        expr = " AND ".join(expressions)
        return expr, args
    
    def add_record(self, msg_id, rec):
        """Add a new Task Record, by msg_id."""
        d = self._defaults()
        d.update(rec)
        d['msg_id'] = msg_id
        line = self._dict_to_list(d)
        tups = '(%s)'%(','.join(['?']*len(line)))
        self._db.execute("INSERT INTO %s VALUES %s"%(self.table, tups), line)
        self._db.commit()
    
    def get_record(self, msg_id):
        """Get a specific Task Record, by msg_id."""
        cursor = self._db.execute("""SELECT * FROM %s WHERE msg_id==?"""%self.table, (msg_id,))
        line = cursor.fetchone()
        if line is None:
            raise KeyError("No such msg: %r"%msg_id)
        return self._list_to_dict(line)
    
    def update_record(self, msg_id, rec):
        """Update the data in an existing record."""
        query = "UPDATE %s SET "%self.table
        sets = []
        keys = sorted(rec.keys())
        values = []
        for key in keys:
            sets.append('%s = ?'%key)
            values.append(rec[key])
        query += ', '.join(sets)
        query += ' WHERE msg_id == %r'%msg_id
        self._db.execute(query, values)
        self._db.commit()
    
    def drop_record(self, msg_id):
        """Remove a record from the DB."""
        self._db.execute("""DELETE FROM %s WHERE mgs_id==?"""%self.table, (msg_id,))
        self._db.commit()
    
    def drop_matching_records(self, check):
        """Remove a record from the DB."""
        expr,args = self._render_expression(check)
        query = "DELETE FROM %s WHERE %s"%(self.table, expr)
        self._db.execute(query,args)
        self._db.commit()
        
    def find_records(self, check, id_only=False):
        """Find records matching a query dict."""
        req = 'msg_id' if id_only else '*'
        expr,args = self._render_expression(check)
        query = """SELECT %s FROM %s WHERE %s"""%(req, self.table, expr)
        cursor = self._db.execute(query, args)
        matches = cursor.fetchall()
        if id_only:
            return [ m[0] for m in matches ]
        else:
            records = {}
            for line in matches:
                rec = self._list_to_dict(line)
                records[rec['msg_id']] = rec
            return records

__all__ = ['SQLiteDB']