root/trac/hacks/marketplugin/0.9/tracmarket/db.py

import os
import time
import types
from trac.core import Component

# XXX handle OperationalError for locked db, in select, update, and
# insert; also find all standalone selects and handle there

def kw2sql(*args, **kwargs):
    """

        >>> kw2sql(foo='a', bar='b')
        ('foo = %s AND bar = %s', ['a', 'b'])
        >>> kw2sql(', ', foo='a', bar='b')
        ('foo = %s , bar = %s', ['a', 'b'])
        >>> kw2sql('OR', foo='a', bar='b')
        ('foo = %s OR bar = %s', ['a', 'b'])
        >>> kw2sql('OR', foo='a', bar=('>', 'b'))
        ('foo = %s OR bar > %s', ['a', 'b'])
        >>> kw2sql('OR', dict(foo='a'), dict(foo='b'))
        ('(foo = %s) OR (foo = %s)', ['a', 'b'])
        >>> kw2sql('OR', dict(foo='a'), dict(foo='b'), bar=23, baz=42)
        ('(foo = %s) OR (foo = %s) OR baz = %s OR bar = %s', ['a', 'b', 42, 23])
        >>> class T(DbObject): 
        ...     NAME='footable'
        >>> kw2sql(T, foo='a', bar='b')
        ('footable.foo = %s AND footable.bar = %s', ['a', 'b'])

    """
    parms = []
    terms = []
    op = 'AND'
    table_class = None
    if args:
        for arg in args:
            if arg is None:
                continue
            elif type(arg) is type(str()):
                op = arg
            elif type(arg) in (type(list()), type(tuple())):
                s, p = kw2sql(table_class, *arg)
                if s:
                    terms.append("(%s)" % s)
                    parms += p
            elif type(arg) is type(dict()):
                s, p = kw2sql(table_class, **arg)
                if s:
                    terms.append("(%s)" % s)
                    parms += p
            elif issubclass(arg, DbObject):
                table_class = arg
            else:
                assert False, arg
    if kwargs:
        table_name = None
        if table_class:
            table_name = table_class.NAME
        for var, val in kwargs.items():
            eq = '='
            if val is None:
                continue
            if type(val) in (type(list()), type(tuple())):
                eq = val[0]
                val = val[1]
            if isinstance(val, Select):
                if table_name:
                    terms.append('%s.%s %s (%s)' % (table_name, var, eq, val.sql))
                else:
                    terms.append('%s %s (%s)' % (var, eq, val.sql))
                parms += val.parms
            else:
                if table_name:
                    terms.append("%s.%s %s %%s" % (table_name, var, eq))
                else:
                    terms.append("%s %s %%s" % (var, eq))
                parms.append(val)
    op = ' %s ' % op
    sql = op.join(terms)
    return sql, parms

class Select(object):

    # XXX get rid of env
    def __init__(self, env, db, cls, columns=None, append=None, parms=None,
            order=None, limit=None, *args, **kwargs):
        self.env = env
        self.db = db
        self.cls = cls
        self.debug = os.environ.get("SQLDEBUG", None)
        self.parms = []
        sql = "SELECT "
        if type(columns) is type('a'):
            columns = columns.split(', ')
        if not columns:
            columns = '*'
        sql += '%s ' % ', '.join(columns)
        names = []
        if type(cls) not in (type(list()), type(tuple())):
            cls = (cls, )
        for c in cls:
            names.append(c.NAME)
        sql += 'FROM %s ' % ', '.join(names)
        if args or kwargs:
            s, p = kw2sql(cls[0], *args, **kwargs)
            if s:
                sql += "WHERE %s " % s
                self.parms += p
                if append:
                    sql += "AND "
        elif append:
            sql += "WHERE "
        if append:
            sql += " %s " % append
            if parms:
                self.parms += parms
        if order:
            sql += "ORDER BY %s " % order
        if limit:
            sql += "LIMIT %d " % limit
        self.sql = sql

    def cursor(self):
        cursor = self.db.cursor()
        if self.parms:
            if self.debug:
                self.env.log.debug("'%s', %s" % (self.sql, self.parms))
            cursor.execute(self.sql, self.parms)
        else:
            if self.debug:
                self.env.log.debug("'%s'" % (self.sql))
            cursor.execute(self.sql)
        return cursor

    def rows(self):
        cursor = self.cursor()
        start = time.time()
        while True:
            row = cursor.fetchone() 
            if not row: 
                elapsed = time.time() - start
                if elapsed > .1:
                    self.env.log.info(
                        "query took %f seconds: %s" % (elapsed, self.sql))
                raise StopIteration
            yield row
    rows = property(rows)

    def objs(self):
        assert issubclass(self.cls, DbObject)
        for row in self.rows:
            kwargs = {}
            for var, val in zip(row.keys(), row):
                kwargs[var] = val
            obj = self.cls(self.env, self.db, **kwargs)
            yield obj
    objs = property(objs)

    def table_name(self):
        c = self.cls
        if type(c) in (type(list()), type(tuple())):
            c = c[0]
        return c.NAME
    table_name = property(table_name)

class InitDB(Component):

    # IEnvironmentSetupParticipant methods

    def environment_needs_upgrade(self, db):
        cursor = db.cursor()
        # print self.dbobjects
        for cls in self.dbobjects:
            if cls.upgrade_db(self.env, db, noop=True):
                return True
        return False

    def upgrade_environment(self, db):
        for cls in self.dbobjects:
            cls.upgrade_db(self.env, db)


class DbObject(dict):

    def __init__(self, env, db, **kwargs):
        self._env = env
        self._db = db
        self._set_defaults()
        if not kwargs:
            return 
        for var, val in kwargs.items():
            if val is None:
                continue
            self[var] = val

    def __str__(self):
        import pprint
        return pprint.pformat(self)
        # return str(self)

    def _set_defaults(self):
        pass

    def __getattr__(self, name):
        if name in self.COLUMNS:
            return self.get(name, None)
        raise AttributeError, "column %s not in table %s" % (name, self.NAME)

    def chattr(self, **kwargs):
        for name in kwargs:
            if name not in self.COLUMNS:
                raise AttributeError
        super(DbObject, self).update(kwargs)

    def upgrade_db(cls, env, db, noop=False):
        version = cls.VERSION
        upgrade = False
        cursor = db.cursor()
        cursor.execute("SELECT value FROM system "
                       "WHERE name='market_%s_db_version'" % cls.__name__)
        if cursor.rowcount != 1:
            db_version = 0
        else:
            db_version = int(cursor.fetchone()[0])
        if db_version < cls.VERSION:
            upgrade = True
        if noop:
            return upgrade
        for i in range(db_version, version+1):
            method = "_upgrade_to_v%d" % i
            if hasattr(cls, method):
                getattr(cls, method)(env, db)
                env.log.debug("Done upgrading %s to version %d" %
                        (cls.__name__, i))
        cursor.execute("DELETE FROM system "
                       "WHERE name='market_%s_db_version'" % cls.__name__)
        cursor.execute("INSERT INTO system (name, value) "
                       "VALUES ('market_%s_db_version', %%s)" % cls.__name__,
                       (version, ))
        env.log.debug("Db upgraded: %s" % cls.__name__)
    upgrade_db = classmethod(upgrade_db)

    def _upgrade_to_v1(cls, env, db):
        table = cls.TABLE
        cursor = db.cursor()
        for stmt in db.to_sql(table):
            cursor.execute(stmt)
        env.log.debug("Done creating table %s" % table.name)
    _upgrade_to_v1 = classmethod(_upgrade_to_v1)

    def reload(self):
        '''return a duplicate object from db'''
        # XXX assumes id is primary key (or at least unique)
        return self.__class__.selectone(self._env, self._db, id=self.id)

    def select(cls, env, db, append=None, parms=None, **kwargs):
        s = Select(env, db, cls, append=append, parms=parms, **kwargs)
        return s.objs
    select = classmethod(select)

    def selectone(cls, env, db, *args, **kwargs):
        s = cls.select(env, db, *args, **kwargs)
        try:
            obj = s.next()
        except StopIteration:
            obj = None
        return obj
    selectone = classmethod(selectone)

    def _ck_values(self):
        """Check values before update.  Raise ValueError if bad value."""
        pass 

    def _ignore(self, *args):
        for var in self.COLUMNS:
            if var in args:
                continue
            self._assert(var)

    def _assert(self, *args):
        for var in args:
            if self.get(var, None) is None:
                raise ValueError("missing value: %s" % var)

    def insert(self, db=None, cursor=None):
        if not db:
            db = self._db
        if not cursor:
            cursor = db.cursor()
        self._ck_values()
        sql = "INSERT INTO %s " % self.NAME
        colnames = []
        format = []
        parms = []
        for var, val in self.items():
            colnames.append(var)
            format.append('%s')
            parms.append(val)
        sql += "(%s) " % ', '.join(colnames) 
        sql += "VALUES (%s) " % ', '.join(format) 
        # print sql, parms
        cursor.execute(sql, parms)
        if cursor.rowcount != 1:
            return None
        return self

    def update(self, set=None, parms=None, **kwargs):
        db = self._db
        cursor = db.cursor()
        sql = "UPDATE %s SET " % self.NAME
        if not set:
            set = []
        if not parms:
            parms = []
        for var, val in kwargs.items():
            self[var]=val
            set.append('%s=%%s' % var)
            parms.append(val)
        sql += "%s " % ', '.join(set) 
        kw = {}
        # update this record only
        for col in self.TABLE.key:
            kw[col] = self[col]
        where, whereparms = kw2sql(**kw)
        sql += "WHERE " + where
        parms += whereparms
        # print sql, parms
        self._ck_values()
        cursor.execute(sql, parms)
        return cursor.rowcount

    def delete(self):
        db = self._db
        cursor = db.cursor()
        sql = "DELETE FROM %s WHERE " % self.NAME
        parms = []
        where = []
        # delete this record only
        for col in self.TABLE.key:
            where.append("%s=%%s" % col)
            parms.append(self[col])
        sql += "%s " % ' AND '.join(where)
        # print sql, parms
        cursor.execute(sql, parms)
        return cursor.rowcount


def uniq(set):
    u = {}
    for s in set:
        u[s] = 1
    return u.keys()