root/trunk/storage/storepypgsql.py

# Use libpq directly to avoid all of the DB-API overhead.
from pyPgSQL import libpq

import warnings
import threading
import datetime

import dejavu
from dejavu import storage, logic
from dejavu.storage import db


AdapterToPgSQL = db.AdapterToSQL()
AdapterFromPg = db.AdapterFromDB


class StoreIteratorPgSQL(object):
    """Iterator for populating Units from storage."""
    
    def __init__(self, store, unitClass, expr):
        self.store  = store
        self.unitClass = unitClass
        self.expr = expr
        self.sql, self.imperfect = store.select(unitClass, expr)
    
    def units(self):
        s = self.store
        
        res = s.execute(self.sql)
        if res.resultType != libpq.EMPTY_QUERY:
            columns = {}
            for index in xrange(res.nfields):
                columns[res.fname(index)] = (index, res.ftype(index))
            
            for row in xrange(res.ntuples):
                unit = self.unitClass()
                coercer = AdapterFromPg(unit)
                for key in unit.__class__.properties():
                    index, ftype = columns[key]
                    value = res.getvalue(row, index)
                    try:
                        coercer.consume(key, value, ftype)
                    except Exception, x:
                        x.args += (key, ftype, value)
                        raise x
                # If our SQL is imperfect, don't yield it to the
                # caller unless it passes evaluate().
                if (not self.imperfect) or self.expr.evaluate(unit):
                    yield unit
        res.clear()


class PgSQLDecompiler(db.SQLDecompiler):
    
    def dejavu_icontainedby(self, op1, op2):
        if isinstance(op1, db.ConstWrapper):
            # Looking for text in a field. Use ILike (reverse terms).
            return op2 + " ILIKE '%" + op1.strip("'\"") + "%'"
        else:
            # Looking for field in (a, b, c).
            # Force all args to lowercase for case-insensitive comparison.
            atoms = [self.adapter.coerce(x).lower() for x in op2.basevalue]
            return "LOWER(%s) IN (%s)" % (op1, ", ".join(atoms))
    
    def dejavu_istartswith(self, x, y):
        return x + " ILIKE '" + y.strip("'\"") + "%'"
    
    def dejavu_iendswith(self, x, y):
        return x + " ILIKE '%" + y.strip("'\"") + "'"
    
    def dejavu_ieq(self, x, y):
        # ILIKE with no wildcards should behave like ieq.
        return x + " ILIKE " + y
    
    def dejavu_year(self, x):
        return "date_part('year', " + x + ")"



class StorageManagerPgSQL(storage.StorageManager):
    """StoreManager to save and retrieve Units via pyPgSQL 1.35."""
    
    decompiler = PgSQLDecompiler
    createAdapter = db.FieldTypeAdapter()
    
    def __init__(self, name, arena, allOptions={}):
        storage.StorageManager.__init__(self, name, arena, allOptions)
        
        # connstring = (host=h port=p dbname=d user=u password=p options=o tty=t)
        self.connstring = allOptions[u'Connect']
        atoms = self.connstring.split(" ")
        for atom in atoms:
            k, v = atom.split("=", 1)
            setattr(self, k, v)
        self.CreateIfMissing = allOptions.get(u'Create If Missing', '')
        self.threaded = bool(allOptions.get(u'Threaded', '1'))
        self._connection = None
        
        self.prefix = allOptions.get(u'Prefix', u"djv")
        self.reserve_lock = threading.Lock()
    
    def identifier(self, *atoms):
        ident = '"' + ''.join(map(str, atoms)).replace('"', '""') + '"'
        if len(ident) > 63:
            warnings.warn("Identifier is longer than 63 characters. Most "
                          "installations of Postgres are limited to 63. "
                          "See NAMEDATALEN.")
        return ident
    
    def shutdown(self):
        if self.threaded:
            t = threading.currentThread()
            conn = getattr(t, "dejavu_storage_connection", None)
            if conn is not None:
                conn.finish()
        else:
            if self._connection is not None:
                self._connection.finish()
    
    def _get_conn(self):
        try:
            conn = libpq.PQconnectdb(self.connstring)
        except Exception, x:
            if self.CreateIfMissing:
                self.create_database()
                conn = libpq.PQconnectdb(self.connstring)
            else:
                raise
        return conn
    
    def connection(self):
        if self.threaded:
            t = threading.currentThread()
            if not hasattr(t, 'dejavu_storage_connection'):
                t.dejavu_storage_connection = self._get_conn()
            return t.dejavu_storage_connection
        else:
            if self._connection is None:
                self._connection = self._get_conn()
            return self._connection
    
    def _template_conn(self):
        atoms = self.connstring.split(" ")
        tmplconn = ""
        for atom in atoms:
            k, v = atom.split("=", 1)
            if k == 'dbname': v = 'template1'
            tmplconn += "%s=%s " % (k, v)
        return libpq.PQconnectdb(tmplconn)
    
    def create_database(self):
        self.execute('CREATE DATABASE %s' % self.identifier(self.dbname),
                     self._template_conn())
    
    def drop_database(self):
        self.execute("DROP DATABASE %s;" % self.identifier(self.dbname),
                     self._template_conn())
    
    def select(self, unitClass, expr, distinct_fields=None):
        tablename = self.identifier(self.prefix, unitClass.__name__)
        if distinct_fields:
            distinct_fields = [self.identifier(x) for x in distinct_fields]
            sql = (u'SELECT DISTINCT %s FROM %s' %
                   (u', '.join(distinct_fields), tablename))
        else:
            sql = u'SELECT * FROM %s' % tablename
        w, i = self.where(unitClass, expr)
        if len(w) > 0:
            w = u" WHERE " + w
        else:
            w = u""
        sql += w
        return sql, i
    
    def where(self, cls, expr):
        decom = self.decompiler(self.prefix + cls.__name__, expr)
        return decom.code(), decom.imperfect
    
    def execute(self, query, conn=None):
        if conn is None:
            conn = self.connection()
        try:
            return conn.query(query)
        except Exception, x:
            x.args += (query,)
            raise x
    
    def recall(self, cls, expr=None, pairs=None):
        if expr is None:
            expr = logic.Expression(lambda x: True)
        return StoreIteratorPgSQL(self, cls, expr).units()
    
    def reserve(self, unit):
        """reserve(unit). -> Reserve a persistent slot for unit.
        
        Notice in particular that we do not use the auto-number or
        sequence generation capabilities within some databases, etc.
        The ID should be supplied by UnitSequencers via reserve().
        """
        clsname = unit.__class__.__name__
        tblname = self.identifier(self.prefix, clsname)
        id = self.identifier("ID")
        self.reserve_lock.acquire()
        try:
            if unit.ID is None:
                data = []
                res = self.execute(u'SELECT %s FROM %s;' % (id, tblname))
                if res.resultType != libpq.EMPTY_QUERY:
                    data = [res.getvalue(row, 0) for row in xrange(res.ntuples)]
                unit.ID = unit.sequencer.next(data)
            
            self.execute('INSERT INTO %s (%s) VALUES (%s)' %
                         (tblname, id, AdapterToPgSQL.coerce(unit.ID)))
        finally:
            self.reserve_lock.release()
    
    def save(self, unit, forceSave=False):
        """save(unit, forceSave=False) -> Update storage from unit's data."""
        if unit.dirty() or forceSave:
            cls = unit.__class__
            clsname = cls.__name__
            tablename = self.identifier(self.prefix, clsname)
            
            parms = []
            for key in cls.properties():
                if key != "ID":
                    val = AdapterToPgSQL.coerce(getattr(unit, key))
                    parms.append('%s = %s' % (self.identifier(key), val))
            sql = ('UPDATE %s SET %s WHERE %s = %s' %
                   (tablename, u", ".join(parms),
                    self.identifier("ID"),
                    AdapterToPgSQL.coerce(unit.ID, cls.property_type("ID"))))
            self.execute(sql)
            unit.cleanse()
    
    def destroy(self, unit):
        """Delete the unit."""
        # Use a DELETE command instead of a cursor for better performance.
        deleteStatement = (u'DELETE * FROM %s WHERE %s = %s' %
                           (self.identifier(self.prefix, unit.__class__.__name__),
                            self.identifier("ID"),
                            AdapterToPgSQL.coerce(unit.ID)))
        self.execute(deleteStatement)
    
    def create_storage(self, unitClass):
        tblname = self.identifier(self.prefix, unitClass.__name__)
        
        coerce = self.createAdapter.coerce
        fields = []
        for key in unitClass.properties():
            fields.append(u'%s %s' % (self.identifier(key),
                                      coerce(unitClass, key)))
        try:
            self.execute(u'CREATE TABLE %s (%s)' % (tblname, ", ".join(fields)))
        except libpq.OperationalError, x:
            if not x.args[0].endswith(' already exists\n'):
                raise
        else:
            for index in unitClass.indices():
                indexname = self.identifier(self.prefix, "i",
                                            unitClass.__name__, index)
                self.execute(u'CREATE INDEX %s ON %s (%s)'
                             % (indexname, tblname, self.identifier(index)))
    
    def distinct(self, cls, fields, expr=None):
        """Return distinct values for specified fields."""
        if expr is None:
            expr = logic.Expression(lambda x: True)
        
        # ^%$#@! There's no way to handle imperfect queries without
        # creating all involved Units, which defeats the purpose of
        # distinct, which was a speed issue more than anything. Grr.
        sql, imperfect = self.select(cls, expr, fields)
        # Ignore for now.
##        if imperfect:
##            raise ValueError(u"The following query cannot be reliably "
##                             u"returned from a Postgres data source.",
##                             u"distinct()", cls, fields, expr)
        
        res = self.execute(sql)
        if res.resultType == libpq.EMPTY_QUERY:
            return []
        
        coerce = AdapterFromPg().coerce
        data = []
        for row in xrange(res.ntuples):
            coerced_row = []
            for i in xrange(len(fields)):
                expectedType = cls.property_type(field[i])
                actualType = res.ftype(i)
                val = coerce(res.getvalue(row, i), actualType, expectedType)
                coerced_row.append(val)
            data.append(coerced_row)
        return zip(*data)