wizard/sql.py

   1 import sqlalchemy
   2 import os
   3 import pkg_resources
   4 import copy
   5 import decorator
   6
   7 import wizard
   8 from wizard import plugin, shell
   9
  10 # We're going to use sqlalchemy.engine.url.URL as our database
  11 # info intermediate object
  12
  13 def connect(url):
  14     """Convenience method for connecting to a MySQL database."""
  15     engine = sqlalchemy.create_engine(url)
  16     meta = sqlalchemy.MetaData()
  17     meta.bind = engine
  18     meta.reflect()
  19     return meta
  20
  21 def auth(url):
  22     """
  23     If the URL has a database name but no other values, it will
  24     use the global configuration, and then try the database name.
  25
  26     This function implements a plugin interface named
  27     :ref:`wizard.sql.auth`.
  28     """
  29     if not url:
  30         return None
  31     if not url.database:
  32         # it's hopeless
  33         return url
  34     # omitted port and query
  35     if any((url.host, url.username, url.password)):
  36         # don't try for defaults if a few of these were set
  37         return url
  38     for entry in pkg_resources.iter_entry_points("wizard.sql.auth"):
  39         func = entry.load()
  40         r = func(copy.copy(url))
  41         if r is not None:
  42             return r
  43     env_dsn = os.getenv("WIZARD_DSN")
  44     if env_dsn:
  45         old_url = url
  46         url = sqlalchemy.engine.url.make_url(env_dsn)
  47         url.database = old_url.database
  48     return url
  49
  50 def backup(outdir, deployment):
  51     """
  52     Generic database backup function.
  53     """
  54     # XXX: Change this once deployments support multiple dbs
  55     if deployment.application.database == "mysql":
  56         return backup_mysql(outdir, deployment)
  57     else:
  58         raise NotImplementedError
  59
  60 def backup_mysql(outdir, deployment):
  61     """
  62     Database backups for MySQL using the :command:`mysqldump` utility.
  63     """
  64     outfile = os.path.join(outdir, "db.sql")
  65     try:
  66         shell.call("mysqldump", "--compress", "-r", outfile, *get_mysql_args(deployment.dsn))
  67         shell.call("gzip", "--best", outfile)
  68     except shell.CallError as e:
  69         raise BackupDatabaseError(e.stderr)
  70
  71 def restore(backup_dir, deployment):
  72     """
  73     Generic database restoration function.
  74     """
  75     # XXX: see backup
  76     if deployment.application.database == "mysql":
  77         return restore_mysql(backup_dir, deployment)
  78     else:
  79         raise NotImplementedError
  80
  81 def restore_mysql(backup_dir, deployment):
  82     """
  83     Database restoration for MySQL by piping SQL commands into :command:`mysql`.
  84     """
  85     if not os.path.exists(backup_dir):
  86         raise RestoreDatabaseError("Backup %s doesn't exist" % backup_dir.rpartition("/")[2])
  87     sql = open(os.path.join(backup_dir, "db.sql"), 'w+')
  88     shell.call("gunzip", "-c", os.path.join(backup_dir, "db.sql.gz"), stdout=sql)
  89     sql.seek(0)
  90     shell.call("mysql", *get_mysql_args(deployment.dsn), stdin=sql)
  91     sql.close()
  92
  93 def drop(url):
  94     """
  95     Generic drop database function.  Attempts to run ``DROP
  96     DATABASE`` on the database if no plugins succeed.
  97
  98     This function implements the plugin interface named
  99     :ref:`wizard.sql.drop`.
 100     """
 101     r = plugin.hook("wizard.sql.drop", [url])
 102     if r is not None:
 103         return
 104     engine = sqlalchemy.create_engine(url)
 105     engine.execute("DROP DATABASE `%s`" % url.database)
 106
 107 def get_mysql_args(dsn):
 108     """
 109     Extracts arguments that would be passed to the command line mysql utility
 110     from a deployment.
 111     """
 112     args = []
 113     if dsn.host:
 114         args += ["-h", dsn.host]
 115     if dsn.username:
 116         args += ["-u", dsn.username]
 117     if dsn.password:
 118         args += ["-p" + dsn.password]
 119     args += [dsn.database]
 120     return args
 121
 122 class Error(wizard.Error):
 123     """Generic error class for this module."""
 124     pass
 125
 126 class BackupDatabaseError(Error):
 127     """Backup script failed."""
 128     #: String details of failure
 129     details = None
 130     def __init__(self, details):
 131         self.details = details
 132     def __str__(self):
 133         return """
 134
 135 ERROR: Backing up the database failed, details:
 136
 137 %s""" % self.details
 138
 139 class RestoreDatabaseError(Error):
 140     """Restore script failed."""
 141     #: String details of failure
 142     details = None
 143     def __init__(self, details):
 144         self.details = details
 145     def __str__(self):
 146         return """
 147
 148 ERROR: Restoring the database failed, details:
 149
 150 %s""" % self.details
 151
 152 class RemoveDatabaseError(Error):
 153     """Removing the database failed."""
 154     #: String details of failure
 155     details = None
 156     def __init__(self, details):
 157         self.details = details
 158     def __str__(self):
 159         return """
 160
 161 ERROR: Removing the database failed, details:
 162
 163 %s""" % self.details