Add con.select_db() method. (fixes PyMySQL#80)

methane · methane · commit c5916450b676 · 2013-10-01T01:55:52.000+09:00
diff --git a/pymysql/connections.py b/pymysql/connections.py
@@ -706,6 +706,11 @@ def rollback(self):
         self._execute_command(COM_QUERY, "ROLLBACK")
         self._read_ok_packet()
 
+    def select_db(self, db):
+        '''Set current db'''
+        self._execute_command(COM_INIT_DB, db)
+        self._read_ok_packet()
+
     def escape(self, obj):
         ''' Escape whatever value you pass to it  '''
         if isinstance(obj, str_type):
diff --git a/pymysql/tests/test_connection.py b/pymysql/tests/test_connection.py
@@ -41,6 +41,19 @@ def test_autocommit(self):
         cur.execute("SELECT @@AUTOCOMMIT")
         self.assertEqual(cur.fetchone()[0], 0)
 
+    def test_select_db(self):
+        con = self.connections[0]
+        current_db = self.databases[0]['db']
+        other_db = self.databases[1]['db']
+
+        cur = con.cursor()
+        cur.execute('SELECT database()')
+        self.assertEqual(cur.fetchone()[0], current_db)
+
+        con.select_db(other_db)
+        cur.execute('SELECT database()')
+        self.assertEqual(cur.fetchone()[0], other_db)
+
 
 if __name__ == "__main__":
     try:
