adding support for dicts (because this is what MySQLdb *really* uses)

--HG--
extra : convert_revision : svn%3Ab85faa78-f9eb-468e-a121-7cced6da292c%40475

scrapy/trunk/scrapy/tests/test_storedb.py

@@ -1,21 +1,51 @@
+import MySQLdb
+from twisted.trial import unittest
 from datetime import datetime, timedelta
 from time import sleep
-
-from twisted.trial import unittest
+from scrapy.store.db import DomainDataHistory
+from scrapy.utils.db import mysql_connect, parse_uri, URIValidationError
 from scrapy.conf import settings
 
+class ConnectionTestCase(unittest.TestCase):
+    """ Test connection wrapper """
+    test_db = settings.get('TEST_SCRAPING_DB')
+
+    def setUp(self):
+        if not self.test_db:
+            raise unittest.SkipTest("Missing TEST_SCRAPING_DB setting")
+
+        try:
+            mysql_connect(self.test_db)
+        except ImportError:
+            raise unittest.SkipTest("MySQLdb module not available")
+        except MySQLdb.OperationalError:
+            raise unittest.SkipTest("Test database not available at: %s" % self.test_db)
+
+    def test_parse_uri(self):
+        self.assertRaises(URIValidationError, parse_uri, [self.test_db])
+        self.assertRaises(URIValidationError, parse_uri, self.test_db.replace('mysql:', 'gopher:'))
+
+        d = parse_uri(self.test_db)
+
+        self.assertTrue(isinstance(d, dict))
+        self.assertTrue(all(key in d for key in ('user', 'host', 'db')))
+
+    def test_mysql_connect(self):
+        if not self.test_db:
+            raise unittest.SkipTest("Missing TEST_SCRAPING_DB setting")
+
+        self.assertTrue(isinstance(mysql_connect(self.test_db), MySQLdb.connection))
+        self.assertTrue(isinstance(mysql_connect(parse_uri(self.test_db)), MySQLdb.connection))
+
+class ProductComparisonTestCase(unittest.TestCase):
+    """ Test product comparison functions """
+    test_db = settings.get('TEST_SCRAPING_DB')
 
-class ProductComparisonTestCase(unittest.TestCase):
-    """ Test product comparison functions """
-    
     def setUp(self):
-        self.test_db = settings.get('TEST_SCRAPING_DB')
         if not self.test_db:
             raise unittest.SkipTest("Missing TEST_SCRAPING_DB setting")
 
         try:
-            import MySQLdb
-            from scrapy.utils.db import mysql_connect
             mysql_connect(self.test_db)
         except ImportError:
             raise unittest.SkipTest("MySQLdb module not available")
@@ -23,9 +53,6 @@ class ProductComparisonTestCase(unittest.TestCase):
             raise unittest.SkipTest("Test database not available at: %s" % self.test_db)
 
     def test_domaindatahistory(self):
-        from scrapy.store.db import DomainDataHistory
-        from time import sleep
-
         ddh = DomainDataHistory(self.test_db, 'domain_data_history')
         c = ddh.mysql_conn.cursor()
         c.execute("DELETE FROM domain_data_history")
@@ -71,6 +98,6 @@ class ProductComparisonTestCase(unittest.TestCase):
 
         self.assertEqual(list(ddh.getlast_alldomains()),
                         [('scrapy.org', now3, {'surname': 'Doe', 'name': 'John'})])
-        
+
 if __name__ == "__main__":
     unittest.main()

scrapy/trunk/scrapy/utils/db.py

@@ -6,35 +6,50 @@ from scrapy.conf import settings
 from scrapy import log
 from scrapy.core.engine import scrapyengine
 
-def mysql_connect(db_uri, **kwargs):
-    """
-    Connects to a MySQL DB given a mysql URI
-    """
-    import MySQLdb
+class URIValidationError(Exception):
+    pass
+
+def parse_uri(db_uri):
+    """
+    Parse mysql URI and return settings dict
+    """
+    if not (isinstance(db_uri, basestring) and db_uri.startswith('mysql://')):
+        raise URIValidationError("Incorrect MySQL URI: %s" % db_uri)
 
-    if not db_uri or not db_uri.startswith('mysql://'):
-        raise Exception("Incorrect MySQL URI: %s" % db_uri)
     m = re.search(r"mysql:\/\/(?P<user>[^:]+)(:(?P<passwd>[^@]+))?@(?P<host>[^/]+)/(?P<db>.*)$", db_uri)
     if m:
         d = m.groupdict()
         if d['passwd'] is None:
             del(d['passwd'])
+        return d
 
+def mysql_connect(db_uri_or_dict, **kwargs):
+    """
+    Connects to a MySQL DB given a mysql URI
+    """
+    import MySQLdb
+
+    if isinstance(db_uri_or_dict, dict):
+        d = db_uri_or_dict
+    else:
+        d = parse_uri(db_uri_or_dict)
+
+    if d:
         d['charset'] = 'utf8'
         d.update(settings.get("MYSQL_CONNECTION_SETTINGS", {}))
         d.update(kwargs)
-        
+
         dcopy = d.copy()
         if dcopy.get("passwd"):
             dcopy["passwd"] = "********"
         log.msg("Connecting db with settings %s" % dcopy )
-        
+
         conn = MySQLdb.connect(**d)
-        
+
         #this is to maintain active the connection
         def _ping():
             log.msg("Pinging connection to %s/%s" % (d.get('host'), d.get('db')) )
             conn.ping()
         scrapyengine.addtask(_ping, settings.getint("MYSQL_CONNECTION_PING_PERIOD", 600))
-        
+
         return conn