#!/usr/bin/env jython

import re
import sys
import os
import logging

from java.lang import *
from java.sql import *
from xml.etree.ElementTree import ElementTree

'''
Base class of all MySql database classes.
'''
class MySql(object):

    driverName="com.mysql.jdbc.Driver"

    '''
    Constructor
    '''
    def __init__(self, logger, config, dbType, dbName=None):

        # set up logging
        self.logger = logger
        self.config = config
        self.logger.debug("MySql class constructor")

        # open config and grab database parameters
        if not dbName: self.dbName = config.getDbName(dbType)
        else: self.dbName = dbName
        self.dbHost = config.getDbHost(dbType)
        self.dbUser = config.getDbUser(dbType)
        self.dbPass = config.getDbPassword(dbType)

        # a user friendly connection sring
        connectionStr = self.dbName + "@" + self.dbHost

        # set up JDBC connection
        self.url = "jdbc:mysql://"+self.dbHost+"/"+self.dbName+"?autoReconnect=true&user="+self.dbUser+"&password="+self.dbPass
        try:
            self.con = DriverManager.getConnection(self.url)
            self.connectionID = self.getLastConnectionID()
            self.logger.infoPair("Connected to MySQL Db", connectionStr)
        except:
            self.logger.errorPair("Unable to connect to", connectionStr)
            self.connected = False
            raise

        # determine the mysql major version number:
        sql = "SELECT @@version"

        try:
            rs = self.executeQuery(sql)
            rs.first()
            fullVersion = rs.getString(1)
            majVersion = fullVersion[0:3]
            self.version = float(majVersion)
            
        except:
            print "failed to get version number"
            raise

        self.connected = True

    '''
    Disconnect from database
    '''
    def disconnect(self):
        self.con.close()
        self.connected = False

    '''
    Destructor
    '''
    def __del__(self):

        self.logger.debug("MySql destructor")
        self.disconnect()
    
    '''
    Locks a table. This is wait if the locks is already held by another session
    '''
    def lockTable(self, table):
        # print "Locking table: " + table
        self.execute("LOCK TABLES " + table + " WRITE")

    '''
    Unlocks all tables locked by this session.
    '''
    def unlockTables(self):
        # print "Unlocking tables"
        self.execute("UNLOCK TABLES")

    '''
    Kills the last connection ID, so long as it's not THIS connection ID
    '''
    def killLastConnectionID(self):
   
        connectionID = self.getLastConnectionID()
        if connectionID == self.connectionID:
            print "not killing connection: ", connectionID
            self.logger.debug("NOT going to kill THIS connection ID")
            return

        sql = "KILL %d" % connectionID

        try:
            self.execute(sql)
        except:
            print "failed to kill connection: ", connectionID

    def listTables(self):

        sql = "show tables"
        rs = self.executeQuery(sql)  
        while (rs.next()):
            tableName = rs.getString(1)
            # print "table: " + tableName
        rs.close()

    '''
    Checks whether this table exists
    '''
    def tableExists(self, tableName):

        sql = "SELECT COUNT(*) \
               FROM information_schema.tables \
               WHERE table_schema = '" + self.dbName + "' \
               AND table_name = '" + tableName + "'"

        try:
            rs = self.executeQuery(sql)
            rs.first()
            if rs.getInt(1) > 0: return True
            else: return False
        except:
            self.logger.errorPair("Could not check if table exists", sql)
            pass

        return False

    '''
    Checks whether there are any other connections to this database
    '''
    def anyOtherConnections(self):

        sql = "SELECT COUNT(*) \
               FROM INFORMATION_SCHEMA.PROCESSLIST \
               WHERE DB='" + self.dbName + "'"
        try:
            rs = self.executeQuery(sql)
            rs.first()
            if rs.getInt(1) > 1: return True
            else: return False
        except:
            self.logger.errorPair("Could not check for connection", sql)
            pass

        return True

    '''
    Gets the last connection ID
    '''
    def getLastConnectionID(self):

        sql = "SELECT ID \
               FROM INFORMATION_SCHEMA.PROCESSLIST \
               WHERE DB='" + self.dbName + "' \
               ORDER BY ID"
        rs = self.executeQuery(sql)
        rs.last()
        return rs.getInt(1)

    '''
    Updates all rows of this column of this table with this value
    '''
    def updateAllRows(self, table, column, value):

        sql = "UPDATE " + table + " SET " + column + " = " + value
        self.execute(sql)

    '''
    OPTIMIZEs a table
    '''
    def optimizeTable(self, table):

        sql = "OPTIMIZE TABLE " + table
        try: self.execute(sql)
        except: return

    '''
    Creates a MySQL-friendly table name from the provided string
    '''
    def getDbFriendlyTableName(self, origName):

        name = origName
        name = name.replace('.', '_')
        name = name.replace('/', '_')

        return name

    '''
    Drops a column from a table
    '''
    def dropColumn(self, table, column):

        sql = "ALTER TABLE " + table + " DROP " + column
        try: self.execute(sql)
        except: return

    '''
    Add a column to the table
    '''
    def addColumn(self, table, column, type):

        sql = "ALTER TABLE " + table + " ADD COLUMN " + column + " " + type
        try: self.execute(sql)
        except: return

    '''
    Drops a list of tables
    '''
    def dropTables(self, tables):
        for table in tables: self.dropTable(table)

    '''
    Drops a table
    '''
    def dropTable(self, table):

        sql = "DROP TABLE " + table
        try: self.execute(sql)
        except: return

    '''
    Alters a column to be unique
    '''
    def makeColumnUnique(self, table, column):

        self.logger.debug("Making '"+column+"' unique on table '"+table+"'")

        sql = "ALTER TABLE " + table + " ADD UNIQUE (" + column + ")"
        try:
            self.execute(sql)
        except: pass
            #self.logger.warn("Index already in place on '" + column + "' for table '" + table + "'")

    '''
    Alters a column to be the primary key
    '''
    def makeColumnPrimaryKey(self, table, column):

        self.logger.debug("Making '"+column+"' the primary key on table '"+table+"'")

        sql = "ALTER TABLE " + table + " ADD PRIMARY KEY (" + column + ")"
        try:
            self.execute(sql)
        except: pass
            #self.logger.warn("Index already in place on '" + column + "' for table '" + table + "'")
    '''
    Changes the database engine for this table to InnoDb
    '''
    def changeEngineToInnoDB(self, table):
        self.changeEngine(table, "InnoDB")

    '''
    Changes the database engine for this table
    '''
    def changeEngine(self, table, engine):

        sql = "ALTER TABLE " + table + " ENGINE=" + engine
        try: 
            self.execute(sql)
        except: 
            self.logger.errorPair("Could not change table engine", sql)
            return False

        return True

    '''
    Adds an index to the supplied table and column
    '''
    def createIndex(self, table, column):

        #self.logger.debug("Creating index on column '"+column+"' for table '"+table+"'")

        sql = "CREATE INDEX "+table+"_"+column+"_index ON "+table+" ("+column+")"
        try:
            self.execute(sql)
        except: pass
            #self.logger.warn("Index already in place on '" + column + "' for table '" + table + "'")

    def addPrimaryKeyWithIgnore(self, table, column):
        # note that in mysql versions later than 5.1, this fails
        # unless the following is called first: 
        # "set session old_alter_table=1"

        # follow the command with 
        # "set session old_alter_table=0"

        # OF COURSE, this fails for mysql version <= 5.1...
        if self.scratchDb.version > 5.1:
            self.scratchDb.execute("set session old_alter_table=1")

        # I should probably do this before the whole table is built, right?
        sql = "ALTER IGNORE TABLE " + table + " ADD PRIMARY KEY (" + column + ")"
        self.scratchDb.execute(sql)
        
        if self.scratchDb.version > 5.1:
            self.scratchDb.execute("set session old_alter_table=0")

    '''
    Checks that we have a connection
    '''
    def isConnected(self):

        try: 
            self.con
            self.connected = True
        except:
            self.logger.errorPair("Not connected to", self.url)
            self.connected = False

        return self.connected

    '''
    TODO
    '''
    def execute(self, sql):

        if not self.isConnected(): raise

        stmt = self.con.createStatement()
        stmt.execute(sql)
        stmt.close()

    '''
    TODO
    '''
    def executeQuery(self, sql):

        if not self.isConnected(): raise
  
        stmt = self.con.createStatement()
        rs = stmt.executeQuery(sql)
        #stmt.close()
        return rs

    '''
    TODO
    '''
    def executeUpdatableQuery(self, sql):

        stmt = self.con.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE)
        rs = stmt.executeQuery(sql)
        #stmt.close()
        return rs
    '''
    Returns a list of column names for this table
    '''
    def getColumnNames(self, tableName):

       sql = "SHOW COLUMNS FROM " + tableName
       rs = self.executeQuery(sql)
       columns = []
       while (rs.next()): columns.append(rs.getString(1))
       rs.close()
       
       return columns

    '''
    Replaces all NULL values in the provided table and column ith the provided substitute 
    '''
    def replaceNullsInThisColumn(self, tableName, column, sub):

      sql = "UPDATE " + tableName + " SET " + column + " = " + sub + " WHERE " + column + " IS NULL"
      self.execute(sql)

    '''
    Replaces all NULL values in the provided table with the provided substitute 
    '''
    def replaceNulls(self, tableName, sub):

       # get list of columns
       columns = self.getColumnNames(tableName)

       # now loop through all columns and replace all NULLs with sub
       for column in columns:
          
          sql = "UPDATE " + tableName + " SET " + column + " = " + sub + " WHERE " + column + " IS NULL"
          self.execute(sql)

    '''
    Reports the number of rows with this columns set as this numeric value and then deletes them 
    '''
    def reportAndDeleteRowsWithThisValue(self, tableName, columnName, value):

        sql = "SELECT COUNT(*) FROM " + tableName + " WHERE " + columnName + " = " + value
        rs = self.executeQuery(sql)
        rs.first()
        nBad = rs.getInt(1)

        sql="DELETE FROM " + tableName + " WHERE " + columnName + " = " + value
        self.execute(sql)
        self.logger.debugPair("%s rows with %s" % (columnName, value), "%d deleted" % nBad)

        return nBad

    '''
    Reports the number of rows with this column NULL and then deletes them 
    '''
    def reportAndDeleteRowsWithNULLS(self, tableName, columnName):

        sql = "SELECT COUNT(*) FROM " + tableName + " WHERE " + columnName + " IS NULL"
        rs = self.executeQuery(sql)
        rs.first()
        nBad = rs.getInt(1)

        sql="DELETE FROM " + tableName + " WHERE " + columnName + " IS NULL"
        self.execute(sql)
        self.logger.debugPair("NULL %s values in %s" % (columnName, tableName), "%d deleted" % nBad)

        return nBad

    '''
    Searches a table and reports the columns that are either partially or completely populated with NULLs
    '''
    def reportNulls(self, tableName, showPartials):

       # first, count rows
       sql = "SELECT COUNT(*) FROM " + tableName
       rs = self.executeQuery(sql)
       rs.first()
       numRows = rs.getInt(1)

       # get list of columns
       columns = self.getColumnNames(tableName)

       print "+----------------------+---------------+"
       print "|  %25s           |" % tableName
       print "+----------------------+---------------+"

       # now see which columns are full of NULLS, with are partially NULL
       for column in columns:
          
          sql = "SELECT COUNT(*) FROM " + tableName + " WHERE " + column + " IS NULL"
          rs = self.executeQuery(sql)
          rs.first()
          if rs.getInt(1) == numRows:
              print "| %20s | all NULL      |" % column
          elif showPartials and rs.getInt(1) > 0:
              print "| %20s | partial NULL  |" % column
       rs.close()
       print "+----------------------+---------------+"

    '''
    Returns a row count for this table
    '''
    def getRowCount(self, table):

        sql = "SELECT COUNT(*) FROM " + table
        try:
            rs = self.executeQuery(sql)  
            rs.first()
            return rs.getInt(1)
        except:
            self.logger.exception("Could not count rows for table: '" + table + "'")
            return -1

    '''
    Adds a column to this table which contains a row count
    '''
    def addRowCountColumn(self, table, colName):

        self.execute("ALTER TABLE " + table + " ADD COLUMN " + colName + " INTEGER FIRST")
        self.execute("SET @i = 0")
        self.execute("UPDATE " + table + " SET " + colName + "=(@i:=@i+1)")
        self.execute("ALTER TABLE " + table + " ADD PRIMARY KEY (" + colName + ")")

