~m4v/scratbot-plugins/spaces

# -*- Encoding: UTF-8 -*-
###
# Copyright (c) 2008-2009, Elián Hanisch
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
###

#import supybot.log as log

class DatabaseError(Exception):
    pass


class SqlString(str):
    """Because the normal __repr__ screws up
    the encodings when I only want quotes."""
    def __new__(cls, s=''):
        s = s.lower() # always lowercase
        if 'Ñ' in s:
            s = s.replace('Ñ', 'ñ') # need this for spanish
        return str.__new__(cls, s)

    def __repr__(self):
        return '\'%s\'' % self

class Table(object):
    """
    Class for handling a sql table. <tableStruct> should be a tuple with pairs (colunm_name,
    colunm_type) that defines the table  structure. <name> is the table name and <tableIndex> the
    name of the colunm that should be indexed.
    Allows to iterate the structure, and table[<column_nam>] or table[<column_number>] will return the column data type.
    You can get a row of the table with table.row(<tuple_of_values>)."""

    def __init__(self, tableStruct=(), name='', tableIndex=None):
        self.name = name or getattr(tableStruct, 'name', name)
        self.tableIndex = tableIndex or getattr(tableStruct, 'tableIndex', tableIndex)
        if isinstance(tableStruct, Table):
            self._struct = tableStruct._struct
        else:
            self._struct = tableStruct
        self.row = self.makeRow(self._struct)

    @staticmethod
    def makeRow(table):
        """Row class generator for use with this table."""
        class TableRow(Row):
            def __init__(self, values, slice=None):
                return Row.__init__(self, table, values, slice=slice)
        return TableRow

    def sql(self):
        """Returns name in a format suitable for SQL syntax."""
        assert self.name
        return SqlString(self.name)

    def __getitem__(self, keyword):
        """Implements self[<column_name>] or self[<column_number>]."""
        if isinstance(keyword, int):
            return self.values()[keyword]
        return self.values()[self._index(keyword)]

    def _index(self, keyword):
        """Gets the index of a column by name."""
        try:
            return self.keys().index(keyword)
        except IndexError:
            raise KeyError(keyword)

    def keys(self):
        """Returns column names."""
        if not hasattr(self, '_k'):
            self._k = [ pair[0] for pair in self._struct ]
        return self._k

    def values(self):
        """Returns column data types."""
        if not hasattr(self, '_v'):
            self._v = [ pair[1] for pair in self._struct ]
        return self._v

    def items(self):
        """Returns table structure."""
        return self._struct

    def getName(self):
        """Returns table name."""
        return self.name

    def __str__(self):
        return "<Table('%s' struct=%s)>" % (self.name,
                ', '.join(['%s %s' %pair for pair in self._struct]))

    def __iter__(self):
        return iter(self.items())

    def __len__(self):
        return len(self.items())


class ChannelTable(Table):
    """
    This class is for tables associated with a channel. The table name would be name_#channel.
    _globalChannel would be a generic channel name for the case where a specific channel isn't
    defined and means 'any channel', normally just '#'."""

    _globalChannel = '#'
    _channelPrefix = '#'
    def __init__(self, tableStruct=(), name=None, tableIndex=None, channel=None):
        self.__parent = super(ChannelTable, self)
        self.__parent.__init__(tableStruct, name, tableIndex)
        if channel and channel != self._globalChannel:
            self.name = '%s_%s' % (self.getName(), channel)

    def getChannel(self):
        """Returns the channel associated with the table."""
        if self._channelPrefix in self.name:
            return self.name[self.name.find(self._channelPrefix):]
        else:
            return self._globalChannel

    def getName(self):
        """Returns table name."""
        if self._channelPrefix in self.name:
            return self.name[:self.name.find(self._channelPrefix)-1]
        else:
            return self.name

class Row(Table):
    """
    Class for handling rows in a table. This class is particulary useful as it allows to fetch a value
    from a row by the column name: row[<column_name>] instead of using an index in the tuple
    returned by the fetch command. This allows to modify the table structure and maintain the same
    code, and improves code legibility.
    A Row needs the table structure in <tableStruct>, and a subrow can be created by giving a list of
    column names in <slice>"""

    def __init__(self, tableStruct, values, slice=None):
        self.name = ''
        if slice is None:
            assert len(tableStruct) == len(values), \
                    'Values count doesn\'t match, expected %s got %s' % (len(tableStruct), len(values))
            self._values = list(values)
            self._struct = tableStruct
        else:
            # create a subrow with the columns in <slice>
            self._values = []
            self._struct = []
            same_len = len(values) == len(tableStruct)
            if isinstance(tableStruct, Table):
                table = tableStruct
            else:
                table = Table(tableStruct)
            keys = table.keys()
            index = table._index
            append_v = self._values.append
            append_s = self._struct.append
            if not same_len:
                next = iter(values).next
            for key in slice:
                if key in keys:
                    i = index(key)
                    append_s(table[i])
                if same_len:
                    append_v(values[i])
                else:
                    append_v(next())

    def __setitem__(self, keyword, value):
        """Implement changing a value by row[<column_name>] = <value>."""
        if isinstance(keyword, int):
            self.values()[keyword] = value
        else:
            self.values()[self._index(keyword)] = value

    def values(self):
        return self._values

    def items(self):
        if not hasattr(self, '_i'):
            self._i = zip(self.keys(), self.values())
        return self._i

    def sql(self):
        """Returns a Row suitable for using in SQL syntax"""
        L = []
        iterdatatype = iter([ pair[1] for pair in self._struct ])
        append = L.append
        next = iterdatatype.next
        for v in self.values():
            datatype = next()
            if v == None:
                append('NULL')
            elif datatype[:4] == 'TEXT' or datatype[:4] == 'CHAR' or datatype[:7] == 'VARCHAR':
                v = str(v)
                if '\'' in v:
                    v = v.replace('\'', '\'\'')
                append('\'%s\'' % v)
            else:
                append(str(v))
        return Row(self._struct, L)

    def __str__(self):
        return "<Row%s>" %', '.join([ '%s %s' %pair for pair in self.items() ])

    def __iter__(self):
        return iter(self.values())

# class for manipulate sqlite databases
class SqliteDB(object):
    """Class for manipulate sqlite databases. Used tables are cached and it can create tables when
    needed."""
    def __init__(self, filename):
        import os
        self.filename = os.path.expandvars(filename)
        self.tables = {} # table cache

    def close(self):
        try:
            self.db.close()
            del self.db
        except:
            pass

    def commit(self, cmd):
        """Commit a change in the database."""
        cursor = self.db.cursor()
        #print 'commit:', cmd
        try:
            cursor.execute(cmd)
            self.db.commit()
        except Exception, e:
            self.db.rollback()
            raise DatabaseError, e
        finally:
            cursor.close()

    def fetch(self, cmd):
        """Fetch data from database."""
        cursor = self.db.cursor()
        #print 'fetch:', cmd
        try:
            cursor.execute(cmd)
            fetch = cursor.fetchall()
        except Exception, e:
            self.db.rollback()
            raise DatabaseError, e
        finally:
            cursor.close()
        return fetch

    def setDB(self):
        """Connects to the database."""
        try:
            return self.db
        except AttributeError:
            try:
                import sqlite
            except ImportError:
                raise DatabaseError, 'You need to have PySQLite installed to use this plugin.'
        db = sqlite.connect(self.filename, encoding='utf-8')
        self.db = db
        return db

    def createTable(self, table):
        """Creates a table, as defined in <table>. It must be a Table object."""
        fields = ', '.join([ '%s %s' % pair for pair in table ])
        sqltable = table.sql()
        cmd = "CREATE TABLE %r ( %s )" % (sqltable, fields)
        self.commit(cmd)
        if table.tableIndex:
            cmd = "CREATE UNIQUE INDEX 'idx_%s' ON %r (%s)" % (sqltable, sqltable, table.tableIndex)
            self.commit(cmd)

    def checkTable(self, table):
        """Checks the table in the database is how we expect it to be. Raises exception if not."""
        try:
            cmd = "SELECT sql FROM sqlite_master WHERE name=%r" % table.sql()
            s = self.fetch(cmd)[0][0]
        except Exception, e:
            return False
        found = s[s.find('('):][2:-2]
        expect = ', '.join([ '%s %s' % pair for pair in table ])
        if expect != found:
            raise DatabaseError, 'La tabla %s en la base no es la esperada.' % table.name
        return True

    def setTable(self, table, create=False):
        """Sets <table> and returns Table object if the table exists and is valid, otherwise returns
        None. If <create> is True and the table isn't found it will be created.
        <table> must be a Table object."""
        assert isinstance(table, Table), 'table %s isn\'t a Table instance' %str(table)
        tblName = table.sql()
        if tblName in self.tables: # check if in cache
            self.table = self.tables[tblName]
            return self.table
        if self.checkTable(table): # check if the table exists and if is valid
            self.tables[tblName] = table
            self.table = table
            return self.table
        if not create:
            # no table and don't create it
            self.table = None
            return self.table
        self.createTable(table)
        self.tables[tblName] = table
        self.table = table
        return self.table

    def getTables(self):
        """Returns a list of the table names in the database."""
        cmd = "SELECT type,tbl_name FROM sqlite_master"
        tables = [ tbl[1] for tbl in self.fetch(cmd) if tbl[0] == 'table' ]
        if tables:
            return tables
        return ()

    def setup(self, table=None, **kwargs):
        """Shortcut for set the database and table."""
        assert table or self.table
        if not table:
            table = self.table
        self.setDB()
        self.setTable(table, **kwargs)
        return self.table

    def _insert(self, values, table=None, replace=False):
        """Inserts an entry into <table>."""
        if not table:
            table = self.table
        fields = ', '.join(values.sql())
        if replace:
            cmd = "INSERT OR REPLACE INTO %r VALUES ( %s )" % (table.sql(), fields)
        else:
            cmd = "INSERT INTO %r VALUES ( %s )" % (table.sql(), fields)
        self.commit(cmd)

    def add(self, values, table):
        """Adds an entry into <table>, creates table if necessary."""
        self.setup(table, create=True)
        self._insert(values)

    def delete(self, where, equal, table):
        """Deletes an matching entry from <table>."""
        if self.setup(table):
            cmd = "DELETE FROM %r WHERE %s=%r" % (self.table.sql(), where, equal)
            self.commit(cmd)

    def update(self, values, match=None, where='', equal='', table=None, **kwargs):
        """Updates an entry in <table>."""
        if self.setup(table):
            if match is not None:
                where = match
                equal = values[match]
            cmd = self._makeUPDATE(values, where=where, equal=equal, **kwargs)
            if cmd:
                self.commit(cmd)

    def _get(self, table=None, **kwargs):
        """Runs a SELECT query in <table> and returns the result."""
        #print '_get', kwargs
        if self.setup(table):
            cmd = self._makeSELECT(**kwargs)
            return self.fetch(cmd)
        return ()

    def get(self, select='', **kwargs):
        """Returns Row objects of the result of a SELECT query."""
        #print 'get', kwargs
        fetch = self._get(select=select, **kwargs)
        if fetch:
            if not select:
                return map(self.table.row, fetch)
            else:
                slice = select.split(',')
                return [ self.table.row(v, slice=slice) for v in fetch ]
        return ()

    def getOne(self, **kwargs):
        """Returns one Row object."""
        #print 'getOne', kwargs
        fetch = self.get(limit=1, **kwargs)
        if fetch:
            return fetch[0]

    @staticmethod
    def _makePattern(s):
        """Converts the usual * ? pattern match to SQL syntax."""
        if '*' in s:
            s = s.replace('*', '%')
        if '?' in s:
            s = s.replace('?', '_')
        return s

    def _makeUPDATE(self, values, table=None, where='', equal='', like='', last_values=None):
        """Creates an UPDATE query, <last_values> are used for compare and only update fields that
        changed."""
        assert values
        if not table:
            table = self.table
        cmds = [ "UPDATE %r" % table.sql() ]
        append = cmds.append
        if last_values and values[where] == last_values[where]:
            values = values.sql()
            last_values = last_values.sql()
            set = ', '.join([ '%s=%s' % (k, values[k]) for k in values.keys() if values[k] != last_values[k] ])
            if not set:
                # nothing to update
                return ''
        else:
            values = values.sql()
            set = ', '.join([ '%s=%s' % (k, values[k]) for k in values.keys() ])
        append("SET %s" % set)
        if where:
            if equal:
                append("WHERE %s=%r" % (where, equal))
            elif like:
                like = self._makePattern(like)
                append("WHERE %s LIKE %r" % (where, like))
            else:
                append("WHERE %s" %where)
        return ' '.join(cmds)

    def _makeSELECT(self, table=None, select='', where='', equal='', like='', order='', limit=None):
        """Creates a SELECT query."""
        if not table:
            table = self.table
        if not select:
            select = '*'
        cmds = [ "SELECT %s FROM %r" %(select, table.sql()) ]
        append = cmds.append
        if where:
            if equal:
                append("WHERE %s=%r" % (where, equal))
            elif like:
                like = self._makePattern(like)
                append("WHERE %s LIKE %r" % (where, like))
            else:
                append("WHERE %s" %where)
        if order:
            append("ORDER BY %s" % order)
        if limit is not None:
            append("LIMIT %s" % limit)
        return ' '.join(cmds)


class ChannelSqliteDB(SqliteDB):
    """
    For databases that need ChannelTable tables."""
    def __init__(self, *args):
        self.__parent = super(ChannelSqliteDB, self)
        self.__parent.__init__(*args)

    def setTable(self, table, channel=None, **kwargs):
        if channel is not None:
            table = ChannelTable(table, channel=channel)
        return self.__parent.setTable(table, **kwargs)

    def add(self, values, table, channel=None):
        self.setup(table, channel=channel, create=True)
        self._insert(values)

    def _get(self, table=None, channel=None, **kwargs):
        if self.setup(table, channel=channel):
            cmd = self._makeSELECT(**kwargs)
            return self.fetch(cmd)
        return ()


###################
### MySQL stuff ###
# XXX this is old and broken, and since I'm not interested in MySQL atm it's commented out
#
#class MyTable(Table):
#   """Special table for MySQL, table names should be quoted with `` in mysql."""
#   def __repr__(self):
#       return '`%s`' % self.name
#
#class MysqlDB(SqliteDB):
#   """Class for manipulate MySQL databases."""
#   tableFact = MyTable(tblFactFields, 'factos', 'name(25)')
#   def __init__(self, host, user, passwd, dbname):
#       self.data = (host, user, passwd, dbname)
#
#   def setDB(self):
#       if hasattr(self, 'db'):
#           return
#       try:
#           import MySQLdb as mysql
#           self.ProgError = mysql.ProgrammingError
#       except ImportError:
#           raise Exception, 'You need to have MySQLdb installed to use this plugin with mysql databases.'
#       data = self.data
#       try:
#           db = mysql.connect(host=data[0], user=data[1], passwd=data[2], db=data[3])
#       except Exception, e:
#           raise Exception, e
#       self.db = db
#       self.tables = {}
#
#   def createTable(self, table):
#       fields = ', '.join([ '%s %s' % (k, v) for (k, v) in table ])
#       cmd = "CREATE TABLE %r ( %s ) DEFAULT CHARSET utf8" % (table, fields)
#       self.commit(cmd)
#       if hasattr(table, 'index'):
#            cmd = "CREATE UNIQUE INDEX %r ON %r (%s)" % (table, table, table.index)
#            self.commit(cmd)
#
#   def checkTable(self, table):
#       try:
#           cmd = "DESCRIBE %r" % (table, )
#           fetch = self.fetch(cmd)
#       except self.ProgError, e:
#           raise Exception, e
#       except:
#           return False
#       found = tuple([ f[0] for f in fetch ])
#       expect = table.keys()
#       if expect != found:
#           #print found
#           #print expect
#           raise Exception, 'La tabla \'%s\' no es la esperada' % (table, )
#       return True
#
#
## vim:set shiftwidth=4 softtabstop=4 tabstop=4 expandtab textwidth=100: