From 8dc9253ab880c57cdf3240f30767cd43915603f3 Mon Sep 17 00:00:00 2001
From: Guilherme Chapiewski <guilherme.chapiewski@gmail.com>
Date: Thu, 12 Mar 2009 01:19:24 -0300
Subject: [PATCH] #22: Created option to drop database first.

---
 src/simple_db_migrate/main.py  |  2 +-
 src/simple_db_migrate/mysql.py | 23 +++++++++++++++++------
 tests/mysql_test.py            | 22 +++++++++++++++++++++-
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/src/simple_db_migrate/main.py b/src/simple_db_migrate/main.py
index 1a300be..b3f77e1 100644
--- a/src/simple_db_migrate/main.py
+++ b/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:
diff --git a/src/simple_db_migrate/mysql.py b/src/simple_db_migrate/mysql.py
index 2964a45..7dc8069 100644
--- a/src/simple_db_migrate/mysql.py
+++ b/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)
diff --git a/tests/mysql_test.py b/tests/mysql_test.py
index bfa1e83..ad22aad 100644
--- a/tests/mysql_test.py
+++ b/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()