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