root/trunk/storage/storemysql.py

"""
References the MySQLdb package at:
http://sourceforge.net/projects/mysql-python

From the MySQL manual:

"If the server SQL mode has ANSI_QUOTES enabled, string literals can be
quoted only with single quotes. A string quoted with double quotes will be
interpreted as an identifier."

So use single quotes throughout.
"""

# Use _mysql directly to avoid all of the DB-API overhead.
import _mysql
import warnings
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 import storage, logic
from dejavu.storage import db


class AdapterToMySQL(db.AdapterToSQL):
    
    escapes = [("'", "''"), ("\\", r"\\")]
    like_escapes = [("%", r"\%"), ("_", r"\_")]
    
    # TRUE and FALSE only work with 4.1 or better.
    bool_true = "1"
    bool_false = "0"
    
    def coerce_str(self, value, skip_encoding=False):
        if not skip_encoding and not isinstance(value, str):
            value = value.encode(self.encoding)
        return "'" + _mysql.escape_string(value) + "'"
    
    def coerce_bool(self, value):
        # TRUE and FALSE only work with 4.1 or better.
        if value:
            return '1'
        return '0'


class AdapterFromMySQL(db.AdapterFromDB):
    
    def coerce_bool(self, value, coltype):
        if isinstance(value, basestring):
            # either '0' or '1'
            value = (value == '1')
        return bool(value)


class MySQLDecompiler(db.SQLDecompiler):
    
    def dejavu_today(self):
        return "CURDATE()"


class MySQLDecompiler411(MySQLDecompiler):
    # Before MySQL 4.1.1, BINARY comparisons could use UPPER()
    # or LOWER() to perform case-insensitive comparisons. Newer
    # versions must use CONVERT() to obtain a case-sensitive
    # encoding, like utf8.
    
    def dejavu_icontainedby(self, op1, op2):
        if isinstance(op1, db.ConstWrapper):
            # Looking for text in a field. Use Like (reverse terms).
            return ("CONVERT("+ op2 + " USING utf8) LIKE '%" +
                    self.adapter.escape_like(op1) + "%'")
        else:
            # Looking for field in (a, b, c).
            atoms = [self.adapter.coerce(x) for x in op2.basevalue]
            return "CONVERT(%s USING utf8) IN (%s)" % (op1, ", ".join(atoms))
    
    def dejavu_istartswith(self, x, y):
        return ("CONVERT(" + x + " USING utf8) LIKE '" +
                self.adapter.escape_like(y) + "%'")
    
    def dejavu_iendswith(self, x, y):
        return ("CONVERT(" + x + " USING utf8) LIKE '%" +
                self.adapter.escape_like(y) + "'")
    
    def dejavu_ieq(self, x, y):
        return "CONVERT(" + x + " USING utf8) = " + y


class FieldTypeAdapterMySQL(db.FieldTypeAdapter):
    """Return the SQL typename of a DB column."""
    
    float_max_precision = 53
    numeric_max_precision = 16
    numeric_max_bytes = 8
    
    def float_type(self, precision):
        """Return a datatype which can handle the given precision."""
        if precision <= 23:
            return "FLOAT"
        else:
            return "DOUBLE PRECISION"
    
    def coerce_str(self, cls, key):
        prop = getattr(cls, key)
        bytes = int(prop.hints.get('bytes', '0'))
        if bytes:
            # MySQL VARBINARY/BLOBs will do case-sensitive comparisons.
            # They also won't truncate trailing spaces like VARCHAR does.
            if bytes <= 255:
                return "VARBINARY(%s)" % bytes
            elif bytes < 2 ** 16:
                return "BLOB"
            elif bytes < 2 ** 24:
                return "MEDIUMBLOB"
        return "LONGBLOB"
    
    def coerce_bool(self, cls, key):
        # We could use BOOLEAN, but it wasn't introduced until 4.1.0.
        return "BOOL"
    
    def coerce_datetime_datetime(self, cls, key):
        return "DATETIME"
    
    def coerce_int(self, cls, key):
        prop = getattr(cls, key)
        bytes = int(prop.hints.get('bytes', '4'))
        if bytes == 1:
            return "BOOLEAN"
        elif bytes == 2:
            typename = "SMALLINT"
        elif bytes == 3:
            typename = "MEDIUMINT"
        else:
            typename = "INTEGER"
        if isinstance(cls.sequencer, dejavu.UnitSequencerInteger):
            if key in cls.identifiers:
                typename += " AUTO_INCREMENT"
        return typename



class MySQLIndexSet(db.IndexSet):
    
    def __delitem__(self, key):
        t = self.table
        # MySQL might rename multiple-column indices to "PRIMARY"
        for i in t.sm.tables._get_indices(t.name):
            if i.colname == self[key].colname:
                t.sm.execute('DROP INDEX %s ON %s;' % (i.qname, t.qname))


class MySQLColumnSet(db.ColumnSet):
    
    def _rename(self, oldcol, newname):
        # Override this to do the actual rename at the DB level.
        t = self.table
        t.sm.execute("ALTER TABLE %s CHANGE %s %s %s;" %
                     (t.qname, oldcol.qname,
                      t.sm.quote(newname), oldcol.dbtype))


class MySQLTableSet(db.TableSet):
    
    def __setitem__(self, key, table):
        q = self.sm.quote
        
        fields = []
        pk = []
        for colname, col in table.columns.iteritems():
            qname = col.qname
            dbtype = col.dbtype
            fields.append('%s %s' % (qname, dbtype))
            if colname in table.mysql_identifiers:
                if dbtype.endswith('BLOB') or dbtype == 'TEXT':
                    # MySQL won't allow indexes on a BLOB field
                    # without a specific length.
                    qname = "%s(255)" % qname
                pk.append(qname)
        
        pk = ", ".join(pk)
        if pk:
            pk = ", PRIMARY KEY (%s)" % pk
        
        self.sm.execute('CREATE TABLE %s (%s%s);' %
                        (table.qname, ", ".join(fields), pk))
        
        seq = getattr(table, "mysql_sequencer", None)
        if seq:
            # Wow, what a hack. We have to INSERT a dummy row
            # to set the autoincrement initial value, and we
            # can't delete it until after the CREATE INDEX
            # statements (or the counter will revert).
            colname, initial = seq
            self.sm.execute("INSERT INTO %s (%s) VALUES (%s);"
                            % (table.qname, q(colname), initial - 1))
        
        for k, index in table.columns.indices.iteritems():
            dbtype = table.columns[k].dbtype
            if dbtype.endswith('BLOB') or dbtype == 'TEXT':
                # MySQL won't allow indexes on a BLOB field
                # without a specific length.
                self.sm.execute('CREATE INDEX %s ON %s (%s(255));' %
                                (index.qname, table.qname, q(index.colname)))
            else:
                self.sm.execute('CREATE INDEX %s ON %s (%s);' %
                                (index.qname, table.qname, q(index.colname)))
        
        if seq:
            self.sm.execute("DELETE FROM %s" % table.qname)
        
        dict.__setitem__(self, key, table)
    
    def _get_tables(self, conn=None):
        data, _ = self.sm.fetch("SHOW TABLES FROM %s" %
                                self.sm.quote(self.sm.dbname),
                                conn=conn)
        return [self.sm.tableclass(self.sm, row[0], self.sm.quote(row[0]))
                for row in data]
    
    def _get_columns(self, tablename, conn=None):
        # cols are: Field, Type, Null, Key, Default, Extra.
        # See http://dev.mysql.com/doc/refman/4.1/en/describe.html
        q = self.sm.quote
        data, _ = self.sm.fetch("SHOW COLUMNS FROM %s.%s"
                                % (q(self.sm.dbname), q(tablename)),
                                conn=conn)
        cols = []
        for row in data:
            c = self.sm.columnclass(row[0], self.sm.quote(row[0]),
                                    None, None, row[4])
            
            dbtype = row[1]
            parenpos = dbtype.find("(")
            if parenpos > -1:
                c.hints['bytes'] = dbtype[parenpos+1:-1]
                dbtype = dbtype[:parenpos]
            c.dbtype = dbtype
            
            if dbtype in ('tinyint', 'smallint', 'mediumint', 'int', 'integer'):
                c.type = int
            elif dbtype == 'bigint':
                c.type = long
            elif dbtype in ('float', 'double', 'real'):
                c.type = float
            elif dbtype in ('decimal', 'numeric'):
                c.type = decimal.Decimal
            elif dbtype == 'date':
                c.type = datetime.date
            elif dbtype in ('datetime', 'timestamp'):
                c.type = datetime.datetime
            elif dbtype == 'time':
                c.type = datetime.time
            elif dbtype in ('char', 'varchar', 'binary', 'varbinary',
                            'tinyblob', 'tinytext', 'blob', 'text',
                            'mediumblob', 'mediumtext',
                            'longblob', 'longtext'):
                c.type = str
            cols.append(c)
        return cols
    
    def _get_indices(self, tablename, conn=None):
        indices = []
        try:
            # cols are: Table, Non_unique, Key_name, Seq_in_index, Column_name,
            # Collation, Cardinality, Sub_part, Packed, Null, Index_type, Comment
            q = self.sm.quote
            data, _ = self.fetch("SHOW INDEX FROM %s.%s"
                                 % (q(self.sm.dbname), q(tablename)),
                                 conn=conn)
        except _mysql.ProgrammingError, x:
            if x.args[0] != 1146:
                raise
        else:
            for row in data:
                i = self.sm.indexclass(row[2], self.sm.quote(row[2]),
                                       row[0], row[4], None, not row[1])
                indices.append(i)
        return indices



class StorageManagerMySQL(db.StorageManagerDB):
    """StoreManager to save and retrieve Units via _mysql."""
    
    sql_name_max_length = 64
    # MySQL uses case-sensitive database and table names on Unix, but
    # not on Windows. Use all-lowercase identifiers to work around the
    # problem. "Column names, index names, and column aliases are not
    # case sensitive on any platform."
    # If deployers set lower_case_table_names to 1, it would help.
    sql_name_caseless = True
    
    typeAdapter = FieldTypeAdapterMySQL()
    toAdapter = AdapterToMySQL()
    fromAdapter = AdapterFromMySQL()
    
    tablesetclass = MySQLTableSet
    columnsetclass = MySQLColumnSet
    indexsetclass = MySQLIndexSet
    
    def __init__(self, name, arena, allOptions={}):
        connargs = ["host", "user", "passwd", "db", "port", "unix_socket",
                    "conv", "connect_time", "compress", "named_pipe",
                    "init_command", "read_default_file", "read_default_group",
                    "cursorclass", "client_flag",
                    ]
        self.connargs = dict([(k, v) for k, v in allOptions.iteritems()
                              if k in connargs])
        self.dbname = self.connargs['db']
        
        db.StorageManagerDB.__init__(self, name, arena, allOptions)
        
        self.decompiler = MySQLDecompiler
        # Get the version string from MySQL, to see if we need
        # a different decompiler.
        conn = self._template_conn()
        rowdata, cols = self.fetch("SELECT version();", conn)
        conn.close()
        v = rowdata[0][0]
        self._version = storage.Version(v)
        if self._version > storage.Version("4.1.1"):
            self.decompiler = MySQLDecompiler411
    
    def quote(self, name):
        """Return name, quoted for use in an SQL statement."""
        return '`' + name.replace('`', '``') + '`'
    
    def _get_conn(self):
        try:
            conn = _mysql.connect(**self.connargs)
        except _mysql.OperationalError, x:
            if x.args[0] == 1040:   # Too many connections
                raise db.OutOfConnectionsError
            raise
        return conn
    
    def _template_conn(self):
        tmplconn = self.connargs.copy()
        tmplconn['db'] = 'mysql'
        return _mysql.connect(**tmplconn)
    
    def fetch(self, query, conn=None):
        """fetch(query, conn=None) -> rowdata, columns.
        
        rowdata: a nested list (or tuples), column values within rows.
        columns: a series of 2-tuples (or more). The first tuple value
            will be the column name, the second value will be the column
            type.
        """
        if conn is None:
            conn = self.connection()
        self.execute(query, conn)
        # store_result uses a client-side cursor
        res = conn.store_result()
        return res.fetch_row(0, 0), res.describe()
    
    def destroy(self, unit):
        """destroy(unit). Delete the unit."""
        t = self.tables[unit.__class__.__name__].qname
        self.execute('DELETE FROM %s WHERE %s;' % (t, self.id_clause(unit)))
    
    def version(self):
        return "MySQL Version: %s" % self._version
    
    def _seq_UnitSequencerInteger(self, unit):
        """Reserve a unit using the table's AUTO_INCREMENT field."""
        cls = unit.__class__
        t = self.tables[cls.__name__]
        
        fields = []
        values = []
        for key in cls.properties:
            typename = self.typeAdapter.coerce(cls, key)
            if typename.endswith("AUTO_INCREMENT"):
                # Skip this field, since we're using AUTO_INCREMENT
                continue
            val = self.toAdapter.coerce(getattr(unit, key))
            fields.append(t.columns[key].qname)
            values.append(val)
        
        fields = ", ".join(fields)
        values = ", ".join(values)
        
        conn = self.connection()
        self.execute('INSERT INTO %s (%s) VALUES (%s);' %
                     (t.qname, fields, values), conn)
        
        # Grab the new ID. This is threadsafe because db.reserve has a mutex.
        setattr(unit, cls.identifiers[0], conn.insert_id())
    
    #                               Schemas                               #
    
    def create_database(self):
        # _mysql has create_db and drop_db commands, but they're deprecated.
        sql = 'CREATE DATABASE %s;' % self.quote(self.sql_name(self.dbname))
        conn = self._template_conn()
        self.execute(sql, conn)
        conn.close()
    
    def drop_database(self):
        sql = 'DROP DATABASE %s;' % self.quote(self.sql_name(self.dbname))
        conn = self._template_conn()
        self.execute(sql, conn)
        conn.close()
    
    def create_storage(self, cls):
        """Create storage for the given class."""
        colname = self.column_name
        
        # Make a Table object.
        tablename = self.table_name(cls.__name__)
        t = self.tableclass(self, tablename, self.quote(tablename))
        
        indices = cls.indices()
        fields = []
        for key in cls.properties:
            dbtype = self.typeAdapter.coerce(cls, key)
            prop = cls.property(key)
            cname = colname(cls.__name__, key)
            col = self.columnclass(cname, self.quote(cname), dbtype,
                                   prop.type, prop.default, prop.hints.copy())
            # Use the superclass call to avoid ALTER TABLE.
            dict.__setitem__(t.columns, key, col)
            
            if key in indices:
                iname = self.table_name("i" + cls.__name__ + key)
                i = self.indexclass(iname, self.quote(iname), tablename, cname)
                # Use the superclass call to avoid CREATE INDEX.
                dict.__setitem__(t.columns.indices, key, i)
        
        # Hack to get PRIMARY KEY right. See MySQLTableSet.__setitem__
        t.mysql_identifiers = cls.identifiers
        
        # Hack to get AUTO_INCREMENT right where initial > 1.
        # See MySQLTableSet.__setitem__
        if isinstance(cls.sequencer, dejavu.UnitSequencerInteger):
            i = cls.sequencer.initial
            if i > 1:
                t.mysql_sequencer = (t.columns[cls.identifiers[0]].name, i)
        
        # Attach to self.tables, which should call CREATE TABLE.
        self.tables[cls.__name__] = t