#22: Created option to drop database first.

src/simple_db_migrate/main.py

@@ -12,7 +12,7 @@ def __init__(self, options=None, args=None, mysql=None, db_migrate=None):
 
         self.__mysql = mysql
         if self.__mysql is None:
-            self.__mysql = MySQL(self.__options.db_config_file)
+            self.__mysql = MySQL(db_config_file=self.__options.db_config_file, drop_db_first=self.__options.drop_db_first)
 
         self.__db_migrate = db_migrate
         if self.__db_migrate is None:

src/simple_db_migrate/mysql.py

@@ -4,7 +4,7 @@
 
 class MySQL(object):
 
-    def __init__(self, db_config_file="simple-db-migrate.conf", mysql_driver=MySQLdb):
+    def __init__(self, db_config_file="simple-db-migrate.conf", mysql_driver=MySQLdb, drop_db_first=False):
         self.__cli = CLI()
 
         # read configurations
@@ -21,9 +21,12 @@ def __init__(self, db_config_file="simple-db-migrate.conf", mysql_driver=MySQLdb
         self.__mysql_user__ = USERNAME
         self.__mysql_passwd__ = PASSWORD
         self.__mysql_db__ = DATABASE
-
-        self.__create_database_if_not_exists()
-        self.__create_version_table_if_not_exists()
+
+        if drop_db_first:
+            self._drop_database()
+
+        self._create_database_if_not_exists()
+        self._create_version_table_if_not_exists()
 
     def __mysql_connect(self, connect_using_db_name=True):
         try:
@@ -38,13 +41,21 @@ def __execute(self, sql):
         db = self.__mysql_connect()
         db.query(sql)
         db.close()
+
+    def _drop_database(self):
+        db = self.__mysql_connect(False)
+        try:
+            db.query("drop database %s;" % self.__mysql_db__)
+        except Exception, e:
+            self.__cli.error_and_exit("can't drop database '%s'; database doesn't exist" % self.__mysql_db__)
+        db.close()
 
-    def __create_database_if_not_exists(self):
+    def _create_database_if_not_exists(self):
         db = self.__mysql_connect(False)
         db.query("create database if not exists %s;" % self.__mysql_db__)
         db.close()
 
-    def __create_version_table_if_not_exists(self):
+    def _create_version_table_if_not_exists(self):
         # create version table
         sql = "create table if not exists __db_version__ ( version varchar(20) NOT NULL default \"0\" );"
         self.__execute(sql)

tests/mysql_test.py

@@ -47,7 +47,27 @@ def test_it_should_create_database_and_version_table_on_init_if_not_exists(self)
 
         mysql = MySQL("test.conf", mysql_driver_mock)
 
-    def test_it_should_execute_migration_up_and_remove_from_schema_version(self):
+    def test_it_should_drop_database_on_init_if_its_asked(self):
+        mysql_driver_mock = Mock()
+        db_mock = Mock()
+        cursor_mock = Mock()
+
+        mysql_driver_mock.expects(at_least_once()).method("connect").will(return_value(db_mock))
+
+        db_mock.expects(at_least_once()).method("close")
+        db_mock.expects(once()).method("query").query(eq("drop database migration_test;"))
+        db_mock.expects(once()).method("query").query(eq("create database if not exists migration_test;"))
+        db_mock.expects(once()).method("query").query(eq("create table if not exists __db_version__ ( version varchar(20) NOT NULL default \"0\" );"))
+        db_mock.expects(once()).method("cursor").will(return_value(cursor_mock))
+
+        cursor_mock.expects(once()).method("execute").execute(eq("select count(*) from __db_version__;"))
+        cursor_mock.expects(once()).method("fetchone").will(return_value("0"))
+
+        db_mock.expects(once()).method("query").query(eq("insert into __db_version__ values (\"0\");"))
+
+        mysql = MySQL(db_config_file="test.conf", mysql_driver=mysql_driver_mock, drop_db_first=True)
+
+    def test_it_should_execute_migration_up_and_update_schema_version(self):
         mysql_driver_mock = Mock()
         db_mock = Mock()
         cursor_mock = Mock()