Package Products :: Package ZenUtils :: Module ZenDB
[hide private]
[frames] | no frames]

Source Code for Module Products.ZenUtils.ZenDB

  1  #! /usr/bin/env python 
  2  ############################################################################## 
  3  #  
  4  # Copyright (C) Zenoss, Inc. 2007, 2009, 2011, all rights reserved. 
  5  #  
  6  # This content is made available according to terms specified in 
  7  # License.zenoss under the directory where your Zenoss product is installed. 
  8  #  
  9  ############################################################################## 
 10   
 11   
 12  import logging 
 13  import optparse 
 14  import os 
 15  import subprocess 
 16  import sys 
 17  import copy 
 18   
 19  from config import ConfigFile 
 20   
 21  logging.basicConfig() 
 22  log = logging.getLogger("zen.zendb") 
 23   
24 -class ZenDBError(Exception):
25 - def __init__(self, msg=None):
26 self.msg = msg
27 - def __str__(self):
28 return repr('ZenDBError: %s' % self.msg)
29
30 -class ZenDB(object):
31 requiredParams = ('db-type', 'host', 'port', 'db', 'user') 32
33 - def __init__(self, useDefault, dsn={}, useAdmin=False):
34 # parameter is muteable, use a local copy 35 dsn = copy.deepcopy(dsn) 36 self._db = useDefault 37 if useDefault in ('zep', 'zodb'): 38 dbparams = self._getParamsFromGlobalConf(useDefault) 39 for setting in dbparams: 40 # only override the dsn settings not already specified 41 if not dsn.get(setting): 42 if setting in ('user', 'password') and useAdmin: 43 # determine if global.conf specifies admin settings 44 key = 'admin-' + setting 45 if key in dbparams: 46 dsn[setting] = dbparams[key] 47 else: 48 dsn[setting] = dbparams[setting] 49 50 # check to confirm we have all db params 51 for setting in self.requiredParams: 52 if not dsn.get(setting): 53 raise ZenDBError('Missing a required DB connection setting ' 54 '(%s), and cannot continue. ' % setting) 55 56 self.dbtype = dsn.pop('db-type') 57 if self.dbtype not in ('mysql', 'postgresql'): 58 raise ZenDBError('%s is not a valid database type.' % self.dbtype) 59 log.debug('db type: %s' % self.dbtype) 60 61 self.dbparams = dsn 62 log.debug('connection params: %s' % str(self.dbparams))
63
64 - def _getParamsFromGlobalConf(self, defaultDb):
65 zenhome = os.environ.get('ZENHOME') 66 if not zenhome: 67 raise ZenDBError('No $ZENHOME set. In order to use default ' 68 'configurations, $ZENHOME must point to the ' 69 'Zenoss install.') 70 else: 71 with open(os.path.join(zenhome, 'etc/global.conf'), 'r') as fp: 72 globalConf = ConfigFile(fp) 73 settings = {} 74 for line in globalConf.parse(): 75 if line.setting: 76 key, val = line.setting 77 if key.startswith(defaultDb + '-'): 78 key = key[len(defaultDb)+1:] 79 settings[key] = val 80 return settings
81
82 - def getConnection(self, extraParams={}):
83 if self.dbtype == 'mysql': 84 from Products.ZenUtils.mysql import MySQLdb 85 p = self.dbparams 86 params = {} 87 params['host'] = p['host'] 88 params['port'] = int(p['port']) 89 params['user'] = p['user'] 90 params['passwd'] = p['password'] 91 params['db'] = p['db'] 92 if 'socket' in p: 93 params['unix_socket'] = p['socket'] 94 elif self._db == 'zep': 95 # ZEP does not use unix_socket, its a java app 96 # assume the socket is the same as zodb's if the 97 # hosts are the same 98 zodbZenDB = ZenDB('zodb') 99 if zodbZenDB.dbparams['host'] == params['host']: 100 if 'socket' in zodbZenDB.dbparams: 101 params['unix_socket'] = zodbZenDB.dbparams['socket'] 102 params.update(extraParams) 103 connection = MySQLdb.connect(**params) 104 return connection 105 else: 106 raise NotImplemented("This method does not support %s" % self.dbtype)
107
108 - def dumpSql(self, outfile=None):
109 # output to stdout if nothing passed in, open a file if a string is passed, 110 # or use an open file if that's passed in 111 if outfile is None: 112 outfile = sys.stdout 113 elif isinstance(outfile, basestring): 114 outfile = open(outfile, 'w') 115 if not isinstance(outfile, file): 116 raise ZenDBError('SQL dump output file is invalid. If you passed in a ' 117 'file name, please confirm that its location is ' 118 'writable.') 119 cmd = None 120 env = os.environ.copy() 121 if self.dbtype == 'mysql': 122 # TODO: Handle compression of stream (--compress)? 123 env['MYSQL_PWD'] = self.dbparams.get('password') 124 cmd = ['mysqldump', 125 '--user=%s' % self.dbparams.get('user'), 126 '--host=%s' % self.dbparams.get('host'), 127 '--port=%s' % str(self.dbparams.get('port')), 128 '--single-transaction', 129 self.dbparams.get('db')] 130 elif self.dbtype == 'postgresql': 131 env['PGPASSWORD'] = self.dbparams.get('password') 132 cmd = ['pg_dump', 133 '--username=%s' % self.dbparams.get('user'), 134 '--host=%s' % self.dbparams.get('host'), 135 '--port=%s' % self.dbparams.get('port'), 136 '--format=p', 137 '--no-privileges', 138 '--no-owner', 139 '--create', 140 '--use-set-session-authorization', 141 self.dbparams.get('db')] 142 if cmd: 143 rc = subprocess.Popen(cmd, stdout=outfile, env=env).wait() 144 if rc: 145 raise subprocess.CalledProcessError(rc, cmd)
146
147 - def executeSql(self, sql=None):
148 cmd = None 149 env = os.environ.copy() 150 if self.dbtype == 'mysql': 151 env['MYSQL_PWD'] = self.dbparams.get('password') 152 cmd = ['mysql', 153 '--batch', 154 '--skip-column-names', 155 '--user=%s' % self.dbparams.get('user'), 156 '--host=%s' % self.dbparams.get('host'), 157 '--port=%s' % str(self.dbparams.get('port')), 158 '--database=%s' % self.dbparams.get('db')] 159 elif self.dbtype == 'postgresql': 160 env['PGPASSWORD'] = self.dbparams.get('password') 161 cmd = ['psql', 162 '--quiet', 163 '--tuples-only', 164 '--username=%s' % self.dbparams.get('user'), 165 '--host=%s' % self.dbparams.get('host'), 166 '--port=%s' % self.dbparams.get('port'), 167 self.dbparams.get('db')] 168 if cmd: 169 p = subprocess.Popen(cmd, env=env, 170 stdin=subprocess.PIPE if sql else sys.stdin) 171 try: 172 if sql: 173 p.communicate(sql) 174 rc = p.wait() 175 if rc: 176 raise subprocess.CalledProcessError(rc, cmd) 177 except KeyboardInterrupt: 178 subprocess.call('stty sane', shell=True) 179 p.kill()
180 181 if __name__ == '__main__': 182 parser = optparse.OptionParser() 183 184 # DB connection params 185 parser.add_option('--usedb', dest='usedb', help='Use default connection settings (zodb/zep)') 186 parser.add_option('--useadmin', action='store_true', dest='useadmin', help='Use Admin creds from --usedb') 187 parser.add_option('--dbtype', dest='dbtype', help='Database Type') 188 parser.add_option('--dbhost', dest='dbhost', help='Database Host') 189 parser.add_option('--dbport', type='int', dest='dbport', help='Database Port') 190 parser.add_option('--dbname', dest='dbname', help='Database Name') 191 parser.add_option('--dbuser', dest='dbuser', help='Database User') 192 parser.add_option('--dbpass', dest='dbpass', help='Database Password') 193 194 # Usage options 195 parser.add_option('--dump', action='store_true', dest='dumpdb', help='Dump database') 196 parser.add_option('--dumpfile', dest='dumpfile', help='Output file for database dump (defaults to STDOUT)') 197 parser.add_option('--execsql', dest='execsql', help='SQL to execute (defaults to STDIN)') 198 199 # logging verbosity 200 parser.add_option('-v', '--logseverity', default='INFO', dest='logseverity', help='Logging severity threshold') 201 202 options, args = parser.parse_args() 203 try: 204 loglevel = int(options.logseverity) 205 except ValueError: 206 loglevel = getattr(logging, options.logseverity.upper(), logging.INFO) 207 log.setLevel(loglevel) 208 209 try: 210 zdb = ZenDB( 211 useDefault=options.usedb, 212 dsn = { 213 'db-type': options.dbtype, 214 'host': options.dbhost, 215 'port': options.dbport, 216 'db': options.dbname, 217 'user': options.dbuser, 218 'password': options.dbpass 219 }, 220 useAdmin=options.useadmin, 221 ) 222 223 if options.dumpdb: 224 zdb.dumpSql(options.dumpfile) 225 else: 226 zdb.executeSql(options.execsql) 227 except ZenDBError as e: 228 log.error(e.msg) 229 sys.exit(-1) 230 except subprocess.CalledProcessError as e: 231 log.error('Error executing command: %s' % repr(e.cmd)) 232 sys.exit(e.returncode) 233