DBStats: New db update routines

Replace updatefix.py with the Updater class.  This streamlines some
of the common tasks and groups database updates by release.  Upgrades
from pre 1.1.x are no longer supported.

1 files changed, 258 insertions, 0 deletions

diff --git a/src/lib/Bcfg2/Server/Reports/Updater/Routines.py b/src/lib/Bcfg2/Server/Reports/Updater/Routines.py
new file mode 100644
index 000000000..1d41848e4
--- /dev/null
+++ b/src/lib/Bcfg2/Server/Reports/Updater/Routines.py
@@ -0,0 +1,258 @@
+import logging
+from django.db.models.fields import NOT_PROVIDED
+from django.db import connection, DatabaseError, backend, models
+from django.core.management.color import no_style
+from django.core.management.sql import sql_create
+import django.core.management
+
+import Bcfg2.Server.Reports.settings
+
+logger = logging.getLogger(__name__)
+
+def _quote(value):
+    """
+    Quote a string to use as a table name or column
+    """
+    return backend.DatabaseOperations().quote_name(value)
+
+
+def _rebuild_sqlite_table(model):
+    """Sqlite doesn't support most alter table statments.  This streamlines the
+       rebuild process"""
+    try:
+        cursor = connection.cursor()
+        table_name = model._meta.db_table
+
+        # Build create staement from django
+        model._meta.db_table = "%s_temp" % table_name
+        sql, references = connection.creation.sql_create_model(model, no_style())
+        columns = ",".join([_quote(f.column) \
+                               for f in model._meta.fields])
+
+        # Create a temp table
+        [cursor.execute(s) for s in sql]
+
+        # Fill the table
+        tbl_name = _quote(table_name)
+        tmp_tbl_name = _quote(model._meta.db_table)
+        # Reset this
+        model._meta.db_table = table_name
+        cursor.execute("insert into %s(%s) select %s from %s;" % (
+            tmp_tbl_name,
+            columns,
+            columns,
+            tbl_name))
+        cursor.execute("drop table %s" % tbl_name)
+
+        # Call syncdb to create the table again
+        django.core.management.call_command("syncdb", interactive=False, verbosity=0)
+        # syncdb closes our cursor
+        cursor = connection.cursor()
+        # Repopulate
+        cursor.execute('insert into %s(%s) select %s from %s;' % (tbl_name,
+                                                                  columns,
+                                                                  columns,
+                                                                  tmp_tbl_name))
+        cursor.execute('DROP TABLE %s;' % tmp_tbl_name)
+    except DatabaseError:
+        logger.error("Failed to rebuild sqlite table %s" % table_name, exc_info=1)
+        raise UpdaterError
+
+
+class UpdaterRoutineException(Exception):
+    pass
+
+
+class UpdaterRoutine(object):
+    """Base for routines."""
+    def __init__(self):
+        pass
+
+    def __str__(self):
+        return __name__
+
+    def run(self):
+        """Called to execute the action"""
+        raise UpdaterRoutineException
+
+
+
+class AddColumns(UpdaterRoutine):
+    """
+    Routine to add new columns to an existing model
+    """
+    def __init__(self, model):
+        self.model = model
+        self.model_name = model.__name__
+
+    def __str__(self):
+        return "Add new columns for model %s" % self.model_name
+
+    def run(self):
+        try:
+            cursor = connection.cursor()
+        except DatabaseError:
+            logger.error("Failed to connect to the db")
+            raise UpdaterRoutineException
+
+        try:
+            desc = {}
+            for d in connection.introspection.get_table_description(cursor, 
+                    self.model._meta.db_table):
+                desc[d[0]] = d
+        except DatabaseError:
+            logger.error("Failed to get table description", exc_info=1)
+            raise UpdaterRoutineException
+
+        for field in self.model._meta.fields:
+            if field.column in desc:
+                continue
+            logger.debug("Column %s does not exist yet" % field.column)
+            if field.default == NOT_PROVIDED:
+                logger.error("Cannot add a column with out a default value")
+                raise UpdaterRoutineException
+
+            sql = "ALTER TABLE %s ADD %s %s NOT NULL DEFAULT " % (
+                    _quote(self.model._meta.db_table),
+                    _quote(field.column), field.db_type(), )
+            db_engine = Bcfg2.Server.Reports.settings.DATABASES['default']['ENGINE']
+            if db_engine == 'django.db.backends.sqlite3':
+                sql += _quote(field.default)
+                sql_values = ()
+            else:
+                sql += '%s'
+                sql_values = (field.default, )
+            try:
+                cursor.execute(sql, sql_values)
+                logger.debug("Added column %s to %s" % 
+                        (field.column, self.model._meta.db_table))
+            except DatabaseError:
+                logger.error("Unable to add column %s" % field.column)
+                raise UpdaterRoutineException
+
+
+class RebuildTable(UpdaterRoutine):
+    """
+    Rebuild the table for an existing model.  Use this if field types have changed.
+    """
+    def __init__(self, model, columns):
+        self.model = model
+        self.model_name = model.__name__
+
+        if type(columns) == str:
+            self.columns = [columns]
+        elif type(columns) in (tuple, list):
+            self.columns = columns
+        else:
+            logger.error("Columns must be a str, tuple, or list")
+            raise UpdaterRoutineException
+
+
+    def __str__(self):
+        return "Rebuild columns for model %s" % self.model_name
+
+    def run(self):
+        try:
+            cursor = connection.cursor()
+        except DatabaseError:
+            logger.error("Failed to connect to the db")
+            raise UpdaterRoutineException
+
+        db_engine = Bcfg2.Server.Reports.settings.DATABASES['default']['ENGINE']
+        if db_engine == 'django.db.backends.sqlite3':
+            """ Sqlite is a special case.  Altering columns is not supported. """
+            _rebuild_sqlite_table(self.model)
+            return
+
+        if db_engine == 'django.db.backends.mysql':
+            modify_cmd = 'MODIFY '
+        else:
+            modify_cmd = 'ALTER COLUMN '
+
+        col_strings = []
+        for column in self.columns:
+            col_strings.append("%s %s %s" % ( \
+                modify_cmd,
+                _quote(column),
+                self.model._meta.get_field(column).db_type()
+            ))
+
+        try:
+            cursor.execute('ALTER TABLE %s %s' %
+                (_quote(self.model._meta.db_table), ", ".join(col_strings)))
+        except DatabaseError:
+            logger.debug("Failed modify table %s" % self.model._meta.db_table)
+            raise UpdaterRoutineException
+
+
+
+class RemoveColumns(RebuildTable):
+    """
+    Routine to remove columns from an existing model
+    """
+    def __init__(self, model, columns):
+        super(RemoveColumns, self).__init__(model, columns)
+
+
+    def __str__(self):
+        return "Remove columns from model %s" % self.model_name
+
+    def run(self):
+        try:
+            cursor = connection.cursor()
+        except DatabaseError:
+            logger.error("Failed to connect to the db")
+            raise UpdaterRoutineException
+
+        try:
+            columns = [d[0] for d in connection.introspection.get_table_description(cursor, 
+                    self.model._meta.db_table)]
+        except DatabaseError:
+            logger.error("Failed to get table description", exc_info=1)
+            raise UpdaterRoutineException
+
+        for column in self.columns:
+            if column not in columns:
+                logger.warning("Cannot drop column %s: does not exist" % column)
+                continue
+
+            logger.debug("Dropping column %s" % column)
+
+            db_engine = Bcfg2.Server.Reports.settings.DATABASES['default']['ENGINE']
+            if db_engine == 'django.db.backends.sqlite3':
+                _rebuild_sqlite_table(self.model)
+            else:
+                sql = "alter table %s drop column %s" % \
+                    (_quote(self.model._meta.db_table), _quote(column), )
+                try:
+                    cursor.execute(sql)
+                except DatabaseError:
+                    logger.debug("Failed to drop column %s from %s" % 
+                        (column, self.model._meta.db_table))
+                    raise UpdaterRoutineException
+
+
+class UpdaterCallable(UpdaterRoutine):
+    """Helper for routines.  Basically delays execution"""
+    def __init__(self, fn):
+        self.fn = fn
+        self.args = []
+        self.kwargs = {}
+
+    def __call__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+        return self
+
+    def __str__(self):
+        return self.fn.__name__
+
+    def run(self):
+        self.fn(*self.args, **self.kwargs)
+
+def updatercallable(fn):
+    """Decorator for UpdaterCallable.  Use for any function passed
+       into the fixes list"""
+    return UpdaterCallable(fn)
+
+