X-Git-Url: https://scripts.mit.edu/gitweb/wizard.git/blobdiff_plain/ea778445b501bf99171df5a99298efa4973a8199..319f82896fb658d9bbaf8d01d6c2c7410eb967c7:/wizard/sql.py

diff --git a/wizard/sql.py b/wizard/sql.py
index 548d8a6..e1719fa 100644
--- a/wizard/sql.py
+++ b/wizard/sql.py
@@ -2,8 +2,10 @@ import sqlalchemy
 import os
 import pkg_resources
 import copy
+import decorator
 
-from wizard import shell
+import wizard
+from wizard import plugin, shell
 
 # We're going to use sqlalchemy.engine.url.URL as our database
 # info intermediate object
@@ -36,7 +38,6 @@ def auth(url):
     for entry in pkg_resources.iter_entry_points("wizard.sql.auth"):
         func = entry.load()
         r = func(copy.copy(url))
-        print r
         if r is not None:
             return r
     env_dsn = os.getenv("WIZARD_DSN")
@@ -45,3 +46,118 @@ def auth(url):
         url = sqlalchemy.engine.url.make_url(env_dsn)
         url.database = old_url.database
     return url
+
+def backup(outdir, deployment):
+    """
+    Generic database backup function.
+    """
+    # XXX: Change this once deployments support multiple dbs
+    if deployment.application.database == "mysql":
+        return backup_mysql(outdir, deployment)
+    else:
+        raise NotImplementedError
+
+def backup_mysql(outdir, deployment):
+    """
+    Database backups for MySQL using the :command:`mysqldump` utility.
+    """
+    outfile = os.path.join(outdir, "db.sql")
+    try:
+        shell.call("mysqldump", "--compress", "-r", outfile, *get_mysql_args(deployment.dsn))
+        shell.call("gzip", "--best", outfile)
+    except shell.CallError as e:
+        raise BackupDatabaseError(e.stderr)
+
+def restore(backup_dir, deployment):
+    """
+    Generic database restoration function.
+    """
+    # XXX: see backup
+    if deployment.application.database == "mysql":
+        return restore_mysql(backup_dir, deployment)
+    else:
+        raise NotImplementedError
+
+def restore_mysql(backup_dir, deployment):
+    """
+    Database restoration for MySQL by piping SQL commands into :command:`mysql`.
+    """
+    if not os.path.exists(backup_dir):
+        raise RestoreDatabaseError("Backup %s doesn't exist" % backup_dir.rpartition("/")[2])
+    sql = open(os.path.join(backup_dir, "db.sql"), 'w+')
+    shell.call("gunzip", "-c", os.path.join(backup_dir, "db.sql.gz"), stdout=sql)
+    sql.seek(0)
+    shell.call("mysql", *get_mysql_args(deployment.dsn), stdin=sql)
+    sql.close()
+
+def drop(url):
+    """
+    Generic drop database function.  Attempts to run ``DROP
+    DATABASE`` on the database if no plugins succeed.
+
+    This function implements the plugin interface named
+    :ref:`wizard.sql.drop`.
+    """
+    r = plugin.hook("wizard.sql.drop", [url])
+    if r is not None:
+        return
+    engine = sqlalchemy.create_engine(url)
+    engine.execute("DROP DATABASE `%s`" % url.database)
+
+def get_mysql_args(dsn):
+    """
+    Extracts arguments that would be passed to the command line mysql utility
+    from a deployment.
+    """
+    args = []
+    if dsn.host:
+        args += ["-h", dsn.host]
+    if dsn.username:
+        args += ["-u", dsn.username]
+    if dsn.password:
+        args += ["-p" + dsn.password]
+    args += [dsn.database]
+    return args
+
+class Error(wizard.Error):
+    """Generic error class for this module."""
+    pass
+
+class BackupDatabaseError(Error):
+    """Backup script failed."""
+    #: String details of failure
+    details = None
+    def __init__(self, details):
+        self.details = details
+    def __str__(self):
+        return """
+
+ERROR: Backing up the database failed, details:
+
+%s""" % self.details
+
+class RestoreDatabaseError(Error):
+    """Restore script failed."""
+    #: String details of failure
+    details = None
+    def __init__(self, details):
+        self.details = details
+    def __str__(self):
+        return """
+
+ERROR: Restoring the database failed, details:
+
+%s""" % self.details
+
+class RemoveDatabaseError(Error):
+    """Removing the database failed."""
+    #: String details of failure
+    details = None
+    def __init__(self, details):
+        self.details = details
+    def __str__(self):
+        return """
+
+ERROR: Removing the database failed, details:
+
+%s""" % self.details