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."""
    
    # This was determined through experimentation. Don't change it.
    numeric_max_precision = 253
    
    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 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()
    
    def __init__(self, name, arena, allOptions={}):
        db.StorageManagerDB.__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']
        
        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 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:
            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."""
        self.execute('DELETE FROM %s WHERE %s;' %
                     (self.table_name(unit.__class__.__name__),
                      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__
        clsname = cls.__name__
        tablename = self.table_name(clsname)
        
        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(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_INSERT_ID();")
        setattr(unit, cls.identifiers[0], data[0][0])
    
    #                               Schemas                               #
    
    def create_database(self):
        # _mysql has create_db and drop_db commands, but they're deprecated.
        sql = 'CREATE DATABASE %s;' % self.sql_name(self.dbname)
        conn = self._template_conn()
        self.execute(sql, conn)
        conn.close()
    
    def drop_database(self):
        sql = 'DROP DATABASE %s;' % self.sql_name(self.dbname)
        conn = self._template_conn()
        self.execute(sql, conn)
        conn.close()
    
    def create_storage(self, cls):
        clsname = cls.__name__
        tablename = self.table_name(clsname)
        typename = self.typeAdapter.coerce
        
        fields = []
        pk = []
        for key in cls.properties:
            qname = self.column_name(clsname, key)
            dbtype = typename(cls, key)
            fields.append('%s %s' % (qname, dbtype))
            if key in cls.identifiers:
                if dbtype.endswith('BLOB') or dbtype == 'TEXT':
                    # MySQL won't allow indexes on a BLOB field
                    # without a specific length.
                    qname = "%s(%s)" % (qname, 255)
                pk.append(qname)
        pk = ", ".join(pk)
        if pk:
            pk = ", PRIMARY KEY (%s)" % pk
        self.execute('CREATE TABLE %s (%s%s);'
                     % (tablename, ", ".join(fields), pk))
        
        hasdummy = False
        if isinstance(cls.sequencer, dejavu.UnitSequencerInteger):
            i = cls.sequencer.initial
            if i > 1:
                # Wow, what a hack. We have to create a dummy row
                # to set the autoincrement initial value, and we
                # can't delete it until after the CREATE INDEX
                # statements below (or the counter will revert).
                colname = self.column_name(clsname, cls.identifiers[0])
                self.execute("INSERT INTO %s (%s) VALUES (%s);"
                             % (tablename, colname, i - 1))
                hasdummy = True
        
        for index in cls.indices():
            i = self.table_name("i" + clsname + index)
            
            dbtype = typename(cls, index)
            if dbtype.endswith('BLOB') or dbtype == 'TEXT':
                # MySQL won't allow indexes on a BLOB field
                # without a specific length.
                self.execute('CREATE INDEX %s ON %s (%s(%s));' %
                             (i, tablename,
                              self.column_name(clsname, index), 255))
            else:
                self.execute('CREATE INDEX %s ON %s (%s);' %
                             (i, tablename,
                              self.column_name(clsname, index)))
        
        if hasdummy:
            self.execute("DELETE FROM %s" % tablename)
    
    def rename_property(self, cls, oldname, newname):
        clsname = cls.__name__
        oldcolname = self.column_name(clsname, oldname)
        newcolname = self.column_name(clsname, newname)
        if oldcolname != newcolname:
            self.execute("ALTER TABLE %s CHANGE %s %s %s;" %
                         (self.table_name(clsname), oldcolname, newcolname,
                          self.typeAdapter.coerce(cls, newname)))
    
    def drop_index(self, cls, name):
        # MySQL might rename multiple-column indices to "PRIMARY"
        clsname = cls.__name__
        names = []
        for i in self.get_indices(self.table_name(clsname, quoted=False)):
            if i.name not in names:
                names.append(i.name)
        for n in names:
            self.execute('DROP INDEX %s ON %s;' %
                         (self.sql_name(n), self.table_name(clsname)))
    
    def get_tables(self, conn=None):
        data, _ = self.fetch("SHOW TABLES FROM %s" % self.dbname,
                             conn=conn)
        return [db.Table(row[0]) for row in data]
    
    def get_columns(self, tablename=None, conn=None):
        # cols are: Field, Type, Null, Key, Default, Extra.
        # See http://dev.mysql.com/doc/refman/4.1/en/describe.html
        data, _ = self.fetch("SHOW COLUMNS FROM %s.%s"
                             % (self.dbname, self.sql_name(tablename)),
                             conn=conn)
        cols = []
        for row in data:
            c = db.Column(row[0], None, row[4])
            
            dbtype = row[1]
            parenpos = dbtype.find("(")
            if parenpos > -1:
                c.hints['bytes'] = dbtype[parenpos+1:-1]
                dbtype = dbtype[:parenpos]
            
            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
            data, _ = self.fetch("SHOW INDEX FROM %s.%s"
                                 % (self.dbname, self.sql_name(tablename)),
                                 conn=conn)
        except _mysql.ProgrammingError, x:
            if x.args[0] != 1146:
                raise
        else:
            for row in data:
                indices.append(db.Index(row[2], row[0], row[4], None, not row[1]))
        return indices