root/trunk/storage/storepypgsql.py

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

import datetime

try:
    # Builtin in Python 2.5?
    decimal
except NameError:
    try:
        # Module in Python 2.3, 2.4
        import decimal
    except ImportError:
        pass

import dejavu
from dejavu.storage import db


class AdapterToPgSQL(db.AdapterToSQL):
    
    like_escapes = [("%", r"\\%"), ("_", r"\\_")]


class FieldTypeAdapterPgSQL(db.FieldTypeAdapter):
    
    def coerce_int(self, cls, key):
        prop = getattr(cls, key)
        if isinstance(cls.sequencer, dejavu.UnitSequencerInteger):
            if key in cls.identifiers:
                return ("INTEGER DEFAULT nextval('%s_%s_seq') NOT NULL"
                        % (cls.__name__, key))
        bytes = int(prop.hints.get('bytes', '4'))
        if bytes == 1:
            return "BOOLEAN"
        elif bytes == 2:
            return "SMALLINT"
        else:
            return "INTEGER"


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 '%" + self.adapter.escape_like(op1) + "%'"
        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 '" + self.adapter.escape_like(y) + "%'"
    
    def dejavu_iendswith(self, x, y):
        return x + " ILIKE '%" + self.adapter.escape_like(y) + "'"
    
    def dejavu_ieq(self, x, y):
        # ILIKE with no wildcards should behave like ieq.
        return x + " ILIKE '" + self.adapter.escape_like(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."""
    
    sql_name_max_length = 63
    close_connection_method = 'finish'
    decompiler = PgSQLDecompiler
    toAdapter = AdapterToPgSQL()
    typeAdapter = FieldTypeAdapterPgSQL()
    
    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['Connect']
        atoms = self.connstring.split(" ")
        for atom in atoms:
            k, v = atom.split("=", 1)
            setattr(self, k, v)
    
    def sql_name(self, name, quoted=True):
        name = db.StorageManagerDB.sql_name(self, name, quoted)
        if quoted:
            name = '"' + name.replace('"', '""') + '"'
        return name
    
    def _get_conn(self):
        try:
            return libpq.PQconnectdb(self.connstring)
        except libpq.DatabaseError, x:
            if x.args[0].startswith('could not connect'):
                raise db.OutOfConnectionsError
            raise
    
    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 version(self):
        c = self._template_conn()
        v = c.version
        c.finish()
        return v
    
    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)))
        
        data = [[res.getvalue(row, col) for col in xrange(res.nfields)]
                for row in xrange(res.ntuples)]
        res.clear()
        
        return data, columns
    
    def _seq_UnitSequencerInteger(self, unit):
        """Reserve a unit using the table's SERIAL field."""
        cls = unit.__class__
        clsname = cls.__name__
        tablename = self.table_name(clsname)
        
        fields = []
        values = []
        for key in cls.properties:
            typename = self.typeAdapter.coerce(cls, key)
            if 'nextval' in typename:
                # Skip this field, since we're using a sequencer
                continue
            val = self.toAdapter.coerce(getattr(unit, key))
            fields.append(self.column_name(clsname, key))
            values.append(val)
        
        fields = ", ".join(fields)
        values = ", ".join(values)
        self.execute('INSERT INTO %s (%s) VALUES (%s);' %
                     (str(tablename), fields, values))
        
        # Grab the new ID. This is threadsafe because db.reserve has a mutex.
        data, col_defs = self.fetch("SELECT last_value FROM %s_%s_seq;"
                                    % (clsname, cls.identifiers[0]))
        setattr(unit, cls.identifiers[0], data[0][0])
    
    #                               Schemas                               #
    
    def create_database(self):
        c = self._template_conn()
        self.execute('CREATE DATABASE %s' % self.sql_name(self.dbname), c)
        c.finish()
    
    def drop_database(self):
        c = self._template_conn()
        self.execute("DROP DATABASE %s;" % self.sql_name(self.dbname), c)
        c.finish()
    
    def has_storage(self, cls):
        # For some odd reason, libpq errors if you try to filter by tablename.
        sql = "SELECT tablename FROM pg_tables"
        data, cols = self.fetch(sql)
        return [self.table_name(cls.__name__, quoted=False)] in data
    
    def create_storage(self, cls):
        """Create storage for the given class."""
        clsname = cls.__name__
        tablename = self.table_name(clsname)
        typename = self.typeAdapter.coerce
        
        fields = []
        for key in cls.properties:
            dbtype = typename(cls, key)
            if 'nextval' in dbtype:
                self.execute("CREATE SEQUENCE %s_%s_seq START %s;"
                             % (clsname, key, cls.sequencer.initial))
            fields.append('%s %s' % (self.column_name(clsname, key), dbtype))
        self.execute('CREATE TABLE %s (%s);' % (tablename, ", ".join(fields)))
        
        for index in cls.indices():
            i = self.table_name("i" + clsname + index)
            self.execute('CREATE INDEX %s ON %s (%s);' %
                         (i, tablename, self.column_name(clsname, index)))
    
    def drop_index(self, cls, name):
        clsname = cls.__name__
        for i in self.get_indices(clsname):
            if i.colname == name:
                self.execute('DROP INDEX %s;' % self.sql_name(i.name))
    
    def get_tables(self, conn=None):
        data, _ = self.fetch("SELECT tablename FROM pg_tables WHERE "
                             "schemaname not in ('information_schema', 'pg_catalog')",
                             conn=conn)
        return [db.Table(row[0]) for row in data]
    
    def get_columns(self, tablename=None, conn=None):
        data, _ = self.fetch("SELECT oid FROM pg_class WHERE relname = '%s'"
                             % tablename, conn=conn)
        table_OID = data[0][0]
        sql = ("SELECT attname, atttypid, attnum, attlen "
               "FROM pg_attribute WHERE attrelid = %s" % table_OID)
        data, _ = self.fetch(sql, conn=conn)
        cols = []
        for row in data:
            name = row[0]
            if name in ('tableoid', 'cmax', 'xmax', 'cmin', 'xmin',
                        'oid', 'ctid'):
                # This is a column which PostgreSQL defines automatically
                continue
            
            # Data type
            dbtype, _ = self.fetch("SELECT typname, typlen FROM pg_type "
                                    "WHERE oid = %s" % row[1])
            if dbtype:
                dbtype = dbtype[0][0]
                if dbtype in ('int2', 'int4'):
                    dbtype = int
                elif dbtype == 'bool':
                    dbtype = bool
                elif dbtype == 'int8':
                    dbtype = long
                elif dbtype in ('float4', 'float8', 'money'):
                    dbtype = float
                elif dbtype == 'numeric':
                    dbtype = decimal.Decimal
                elif dbtype == 'date':
                    dbtype = datetime.date
                elif dbtype in ('timestamp', 'timestamptz'):
                    dbtype = datetime.datetime
                elif dbtype in ('time', 'timetz'):
                    dbtype = datetime.time
                elif dbtype in ('char', 'varchar', 'bpchar', 'text'):
                    dbtype = str
            else:
                dbtype = None
            
            # Default value
            default, _ = self.fetch("SELECT adsrc FROM pg_attrdef "
                                    "WHERE adnum = %s AND adrelid = %s"
                                    % (row[2], table_OID))
            if default:
                default = default[0][0]
                if default.startswith("nextval("):
                    default = None
            else:
                default = None
            
            c = db.Column(row[0], dbtype, default)
            
            bytes = row[3]
            if bytes > 0:
                c.hints['bytes'] = bytes
            
            cols.append(c)
        return cols
    
    def get_indices(self, tablename, conn=None):
        # Get the OID of the parent table.
        data, _ = self.fetch("SELECT oid FROM pg_class WHERE relname = '%s'"
                             % tablename, conn=conn)
        if not data:
            return []
        
        table_OID = data[0][0]
        indices = []
        data, _ = self.fetch("SELECT pg_class.relname, indkey, indisprimary, "
                             "indisunique FROM pg_index LEFT JOIN pg_class "
                             "ON pg_index.indexrelid = pg_class.oid WHERE "
                             "pg_index.indrelid = %s" % table_OID, conn=conn)
        for row in data:
            cols = map(int, row[1].split(" "))
            for col in cols:
                d, _ = self.fetch("SELECT attname FROM pg_attribute "
                                  "WHERE attrelid = %s AND attnum = %s"
                                  % (table_OID, col), conn=conn)
                indices.append(db.Index(row[0], tablename, d[0][0],
                                        bool(row[2]), bool(row[3])))
        
        return indices