extended test coverage for the freeze package, thereby fixing some bugs.

dataset/freeze/app.py

@@ -1,7 +1,7 @@
 import logging
 import argparse
 
-from sqlalchemy.exc import ProgrammingError
+from sqlalchemy.exc import ProgrammingError, OperationalError
 from dataset.util import FreezeException
 from dataset.persistence.table import Table
 from dataset.persistence.database import Database
@@ -11,13 +11,16 @@ from dataset.freeze.format import get_serializer
 
 log = logging.getLogger(__name__)
 
-parser = argparse.ArgumentParser(
-    description='Generate static JSON and CSV extracts from a SQL database.',
-    epilog='For further information, please check the documentation.')
-parser.add_argument('config', metavar='CONFIG', type=str,
-                    help='freeze file cofiguration')
-parser.add_argument('--db', default=None,
-                    help='Override the freezefile database URI')
+
+def create_parser():
+    parser = argparse.ArgumentParser(
+        description='Generate static JSON and CSV extracts from a SQL database.',
+        epilog='For further information, please check the documentation.')
+    parser.add_argument('config', metavar='CONFIG', type=str,
+                        help='freeze file cofiguration')
+    parser.add_argument('--db', default=None,
+                        help='Override the freezefile database URI')
+    return parser
 
 
 def freeze(result, format='csv', filename='freeze.csv', fileobj=None,
@@ -102,28 +105,32 @@ def freeze_export(export, result=None):
         serializer_cls = get_serializer(export)
         serializer = serializer_cls(export, query)
         serializer.serialize()
-    except ProgrammingError as pe:
-        raise FreezeException("Invalid query: %s" % pe)
+    except (OperationalError, ProgrammingError) as e:
+        raise FreezeException("Invalid query: %s" % e)
 
 
-def main():
+def freeze_with_config(config, db=None):
+    for export in config.exports:
+        if db is not None:
+            export.data['database'] = db
+        if export.skip:
+            log.info("Skipping: %s", export.name)
+            continue
+        log.info("Running: %s", export.name)
+        freeze_export(export)
+
+
+def main():  # pragma: no cover
     # Set up default logger.
     logging.basicConfig(level=logging.INFO)
 
     try:
+        parser = create_parser()
         args = parser.parse_args()
-        config = Configuration(args.config)
-        for export in config.exports:
-            if args.db is not None:
-                export.data['database'] = args.db
-            if export.skip:
-                log.info("Skipping: %s", export.name)
-                continue
-            log.info("Running: %s", export.name)
-            freeze_export(export)
+        freeze_with_config(Configuration(args.config), args.db)
     except FreezeException as fe:
         log.error(fe)
 
-if __name__ == '__main__':
+if __name__ == '__main__':  # pragma: no cover
     logging.basicConfig(level=logging.DEBUG)
     main()

dataset/freeze/format/common.py

@@ -28,7 +28,7 @@ class Serializer(object):
         self._get_basepath()
 
         if export.get('filename') == '-':
-            export['fileobj'] = sys.stdout
+            export.data['fileobj'] = sys.stdout
         self.fileobj = export.get('fileobj')
 
     def _get_basepath(self):

dataset/freeze/format/fcsv.py

@@ -34,7 +34,7 @@ class CSVSerializer(Serializer):
 
             # handle fileobj that has been passed in:
             if path is not None:
-                if PY3:
+                if PY3:  # pragma: no cover
                     fh = open(path, 'wt', encoding='utf8', newline='')
                 else:
                     fh = open(path, 'wb')
@@ -42,7 +42,7 @@ class CSVSerializer(Serializer):
                 fh = self.fileobj
 
             writer = csv.writer(fh)
-            if PY3:
+            if PY3:  # pragma: no cover
                 writer.writerow(keys)
             else:
                 writer.writerow([k.encode('utf-8') for k in keys])

dataset/freeze/format/fjson.py

@@ -1,14 +1,16 @@
 import json
-from datetime import datetime
+from datetime import datetime, date
 from collections import defaultdict
 
+from six import PY3
+
 from dataset.freeze.format.common import Serializer
 
 
 class JSONEncoder(json.JSONEncoder):
 
     def default(self, obj):
-        if isinstance(obj, datetime):
+        if isinstance(obj, (datetime, date)):
             return obj.isoformat()
 
 
@@ -37,7 +39,13 @@ class JSONSerializer(Serializer):
         for path, result in self.buckets.items():
             result = self.wrap(result)
 
-            fh = open(path, 'wb') if self.fileobj is None else self.fileobj
+            if self.fileobj is None:
+                if PY3:  # pragma: no cover
+                    fh = open(path, 'w', encoding='utf8')
+                else:
+                    fh = open(path, 'wb')
+            else:
+                fh = self.fileobj
 
             data = json.dumps(result,
                               cls=JSONEncoder,

dataset/persistence/database.py

@@ -3,12 +3,7 @@ import threading
 import re
 from sqlalchemy.util import safe_reraise
 
-try:
-    from urllib.parse import urlencode
-    from urllib.parse import parse_qs
-except ImportError:
-    from urllib import urlencode
-    from urlparse import parse_qs
+from six.moves.urllib.parse import urlencode, parse_qs
 
 from sqlalchemy import create_engine
 from sqlalchemy.pool import NullPool

dataset/persistence/util.py

@@ -3,7 +3,7 @@ from inspect import isgenerator
 
 try:
     from collections import OrderedDict
-except ImportError:
+except ImportError:  # pragma: no cover
     from ordereddict import OrderedDict
 
 from sqlalchemy import Integer, UnicodeText, Float, DateTime, Boolean, types, Table, event