root/trunk/storage/storepypgsql.py

# Use libpq directly to avoid all of the DB-API overhead.
from pyPgSQL import libpq
import time
import datetime
import dejavu
from dejavu.storage import db


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(db.StorageManagerDB):
    """StoreManager to save and retrieve Units via pyPgSQL 1.35."""
    
    identifier_length = 63
    close_connection_method = 'finish'
    decompiler = PgSQLDecompiler
    
    def __init__(self, name, arena, allOptions={}):
        db.StorageManagerDB.__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.retry = 5
    
    def _get_conn(self):
        retry = 0
        while True:
            try:
                return libpq.PQconnectdb(self.connstring)
            except libpq.DatabaseError, x:
                msg = x.args[0]
                if msg.endswith('does not exist\n'):
                    if self.CreateIfMissing:
                        self.create_database()
                        return libpq.PQconnectdb(self.connstring)
                elif msg.startswith('could not connect'):
                    retry += 1
                    if retry < self.retry:
                        time.sleep(retry * 0.1)
                        continue
                    x.args += ("Try increasing your Postgres server's max_connections",)
                raise x
    
    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):
        c = self._template_conn()
        self.execute('CREATE DATABASE %s' % self.identifier(self.dbname), c)
        c.finish()
    
    def drop_database(self):
        c = self._template_conn()
        self.execute("DROP DATABASE %s;" % self.identifier(self.dbname), c)
        c.finish()
    
    def fetch(self, query, conn=None):
        """fetch(query, conn=None) -> rowdata, columns."""
        res = self.execute(query, conn)
        
        columns = []
        if res.resultType != libpq.EMPTY_QUERY:
            for index in xrange(res.nfields):
                columns.append((res.fname(index), res.ftype(index)))
        
        def iterator():
            for row in xrange(res.ntuples):
                yield [res.getvalue(row, col) for col in xrange(res.nfields)]
            # This should be more robust--needs a class with a cleanup call.
            res.clear()
        
        return iterator(), columns