Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- # Nose plugin to dump and restore a database, example of writing a Nose plugin
- import logging
- import tempfile
- import os, sys
- import re
- import subprocess
- from nose.plugins import Plugin
- from nose.util import tolist
- ENGINES = ('mysql',)
- engine_re = '(' + '|'.join(ENGINES) + ')'
- r = engine_re + r'://(?:([\w]+)(?::([\w]+))?@)?([\w.-]+)(?::([\d]{1,5}))?'
- regex = re.compile(r)
- def split_dsn(dsn):
- m = regex.match(dsn)
- if not m:
- raise ValueError('Cannot split DSN from %s' % dsn)
- else:
- return dict(zip(('engine','username','password','host','port'),
- m.groups()))
- def make_option_maker(parser, cmdline_prefix="", env_prefix="", dest_prefix=""):
- """Return a function that can be used to add_option with saner syntax.
- Arguments:
- parser: the parser object that will be used by the returned function
- cmdline_prefix: will be used together with the name to build the command
- line option. a suffix of -- and a prefix of - will be added so you
- do not need to include those (--cmdline_prefix-<name>)
- env_prefix: will be used together with the name to specify the environment
- variable from which the option will default (ENV_PREFIX_<NAME>)
- dest_prefix: will be used together with name to specify the variable name
- in which the value of the option will be stored
- """
- env_prefix = env_prefix.upper()
- if cmdline_prefix and not cmdline_prefix.endswith("-"):
- cmdline_prefix += "-"
- if env_prefix and not env_prefix.endswith("_"):
- env_prefix += "_"
- if dest_prefix and not dest_prefix.endswith("_"):
- dest_prefix += "_"
- def option_maker(name, action = "store", help = "", default = None):
- """Use the given name, action and help string to add an option"""
- if default is None:
- _default = os.environ.get(env_prefix + name.upper())
- else:
- _default = default
- parser.add_option("--" + cmdline_prefix + name.lower(),
- action = action, dest = dest_prefix + name.lower(),
- default= _default,
- help = help + " [" + env_prefix + name.upper() + "]")
- return option_maker
- class DatabaseSnapshotPlugin(Plugin):
- """Plugin that takes a snapshot of Daisy before tests start and restores it
- """
- enabled = True
- score = 1
- def options(self, parser, env=None):
- Plugin.options(self, parser, env)
- omaker = make_option_maker(parser, dest_prefix = "dbdump",
- env_prefix = "NOSE_DBDUMP",
- cmdline_prefix = "dbdump")
- omaker('databases', action="append", help="Dump the selected databases")
- omaker('output', help="Path to store the database dump")
- omaker('dsn', help="engine://[user[:password]@]hostname[:port]")
- omaker('keep', action='store_true', help="Don't delete the database dump")
- def configure(self, options, config):
- self.conf = config
- self.databases = []
- if options.dbdump_databases:
- for lst in [tolist(x) for x in options.dbdump_databases]:
- self.databases.extend(lst)
- self.dsn = options.dbdump_dsn
- self.output = options.dbdump_output
- self.keep = options.dbdump_keep
- if not (self.databases and self.dsn):
- self.enabled = False
- return
- if not self.output:
- f,name = tempfile.mkstemp(suffix=".sql", prefix="nose-dbdump-")
- os.close(f)
- self.output = name
- for k,v in split_dsn(self.dsn).items():
- self.__dict__[k] = v
- def begin(self):
- """Take a snapshot of the database"""
- c = "mysqldump %s --databases %s > %s" % (self._get_connection_params(),
- ' '.join(self.databases),
- self.output)
- r = os.system(c)
- assert r == 0, "Got %s for %s" % (r, c)
- def _get_connection_params(self):
- c = ''
- if self.engine == 'mysql':
- if self.username:
- c += ' --user=%s' % self.username
- if self.password:
- c += ' --password=%s' % self.password
- if self.host == 'localhost' and not self.port:
- pass
- else:
- c += ' --host=%s' % self.host
- if self.port:
- c += ' --port=%s' % self.port
- return c
- def finalize(self, result):
- """Restore database to snapshot, if any"""
- c = "mysql %s < %s" % (self._get_connection_params(), self.output)
- #print c
- r = os.system(c)
- assert r == 0, "Got %s for %s" % (r, c)
- if not self.keep:
- print "Removing output from %s" % self.output
- os.remove(self.output)
- else:
- print "Leaving output at %s" % self.output
Add Comment
Please, Sign In to add comment