#!/usr/bin/python

# python DB skeleton
# written by timeless of Digg, Inc.
#
# prerequisites: mysql libraries, sqlalchemy.
#
# go to the bottom of the script for sample code you should modify. the top part of
# this script is all useful even after put into production.

import time as timemodule
import sys

# SQL stuff
from sqlalchemy import *
from sqlalchemy.ext.sessioncontext import SessionContext
from sqlalchemy.ext.assignmapper import assign_mapper

# need to do this after importing sql alchemy, or shit breaks
import logging

# process input arguments
from optparse import OptionParser
parser = OptionParser()
parser.add_option("-n", "--hostname",
                  dest="dbHost", default="",
                  help="Machine running MySQL instance")
parser.add_option("-d", "--dbname",
                  dest="dbName", default="",
                  help="Name of database on MySQL instance")
parser.add_option("-u", "--user",
                  dest="dbUser", default="",
                  help="Username to login as")
parser.add_option("-p", "--password",
                  dest="dbPass", default="",
                  help="Password to use for login")
parser.add_option("-v", "--verbose",
                  dest="verbose", action="store_true", default=False,
                  help="Be verbose")
parser.add_option("-l", "--logfile",
                  dest="logFile", default="",
                  help="Logfile to send STDOUT to (makes script appear quiet)")
parser.add_option("-m", "--multiple",
                  dest="runMultiple", action="store_true", default=False,
                  help="Allow to run multiple copies of this script concurrently")
(options, args) = parser.parse_args()

logLvl=logging.INFO
if (options.verbose == True):
    logLvl=logging.DEBUG
    #logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
    #logging.getLogger('sqlalchemy.orm.unitofwork').setLevel(logging.DEBUG)

# if they want to log to logfile, send STDOUT, STDERR there
if (options.logFile != ""):
    logging.basicConfig(
        filename=logFile,
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logLvl,
        format="%(levelname)s %(asctime)s %(message)s")
else:
    logging.basicConfig(
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logLvl,
        format="%(levelname)s %(asctime)s %(message)s")

if (options.dbName == ""):
    print parser.format_help()
    logging.critical("You must pass in a database to use. Sorry to be pedantic.")
    sys.exit(255)

# for logging
currTime = timemodule.strftime('%Y-%m-%d %H:%M:%S')

logging.info(" ----- Starting %s", (sys.argv[0]))

# ensure no other copy is running
import commands
if (options.runMultiple == False):
    stringNumRunning = commands.getoutput("""ps auxww | grep '%s' | grep -v grep | wc -l""" % (sys.argv[0]))
    numRunning = int(stringNumRunning) - 1

    if (numRunning > 0):
        logging.critical("I see %d other copies of myself running. Kill them first?" % (numRunning))
        sys.exit(255)

dburi = "mysql://%s:%s@%s:3306/%s" % (options.dbUser, options.dbPass, options.dbHost, options.dbName)
engine = create_engine(dburi)

if (options.verbose):
    conn = engine.connect()
else:
    try:
        conn = engine.connect()
    except:
        logging.critical("Failure to connect to database. Run with verbose option to see full traceback.")
        sys.exit(255)

# try to execute some SQL against the DB
def doSql(sql):
    if (options.verbose):
        sqlresult = conn.execute(sql)
        return sqlresult
    else:
        try:
            sqlresult = conn.execute(sql)
            return sqlresult
        except:
            logging.critical("Failed to execute SQL %s on connection. Run with verbose option to see full traceback." % (sql))
            return False

# ----------------------------------------------------------------------------------- #
# -            above can be left alone. below you want to modify.                   - #
# ----------------------------------------------------------------------------------- #

# select user and host from user - typical
logging.info("Selecting from user table (if it exists)")
sqlresult = doSql("select user,host from user")
for row in sqlresult:
    print "%20s: %s" % (row['user'], row['host'])

# processlist - but no column headers
logging.info("The mysql processlist")
sqlresult = doSql("show processlist")
for row in sqlresult:
    print str(row)

# here we iterate through the columns of a single row - we don't
# know beforehand what the columns will be
logging.info("The mysql slave status")
sqlresult = doSql("show slave status")
for row in sqlresult:
    for key in row.keys():
        print "%30s: %s" % (key, row[key])

# a demo of using a dictionary to define what to do on some tables
logging.info("List of tables to do work on")
dontDoTables = {}
dontDoTables['user'] = 1
dontDoTables['db'] = 1
dontDoTables['func'] = 1
dontDoTables['columns_priv'] = 1
sqlresult = doSql("show tables")
for row in sqlresult:
    if (dontDoTables.has_key(row[0])):
        print "NOT doing %s" % (row[0])
    else:
        print "    doing %s" % (row[0])

conn.close()