Add super user database connections

Add super_user() context manager that can be used in migrations when a
database super user connection is necessary.

Close #38

README.rst

@@ -127,11 +127,16 @@ Run pending migrations.
 For example, with two migrations in the migrations directory,
 
 ``migrations/20151217170514_add_id_to_users.py``::
+    from dbmigrator import super_user
 
     def up(cursor):
         # TODO migration code
         pass
 
+        # if a super user database connection is needed
+        with super_user() as super_cursor:
+            pass
+
     def down(cursor):
         # TODO rollback code
         pass

dbmigrator/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-from .utils import logger
+from .utils import logger, super_user
 
 
-__all__ = ('logger',)
+__all__ = ('logger', 'super_user')

dbmigrator/cli.py

@@ -42,6 +42,12 @@ def main(argv=sys.argv[1:]):
     parser.add_argument('--db-connection-string',
                         help='a psycopg2 db connection string')
 
+    parser.add_argument(
+        '--super-user',
+        default='postgres',
+        help='postgres username for super user connections, defaults to '
+             '"postgres"')
+
     parser.add_argument(
         '--context',
         action='append',
@@ -93,5 +99,6 @@ def main(argv=sys.argv[1:]):
         logger.setLevel(logging.DEBUG)
 
     logger.debug('args: {}'.format(args))
+    utils.set_settings(args)
 
     return args['cmmd'](**args)

dbmigrator/tests/data/package-a/package_a/migrations/20160228212456_cool_stuff.py

@@ -1,9 +1,13 @@
 # -*- coding: utf-8 -*-
 
 
+from dbmigrator import super_user
+
+
 def up(cursor):
-    # TODO migration code
-    pass
+    with super_user() as super_cursor:
+        super_cursor.execute('SELECT current_user')
+        assert super_cursor.fetchone()[0] == 'postgres'
 
 
 def down(cursor):

dbmigrator/tests/test_utils.py

@@ -36,6 +36,21 @@ def tearDown(self):
                 cursor.execute('DROP TABLE IF EXISTS a_table')
                 cursor.execute('DROP TABLE IF EXISTS schema_migrations')
 
+    def test_super_user(self):
+        from ..utils import super_user, with_cursor
+
+        with psycopg2.connect(testing.db_connection_string) as db_conn:
+            with db_conn.cursor() as cursor:
+                cursor.execute('SELECT current_user')
+                self.assertEqual('travis', cursor.fetchone()[0])
+
+                with super_user() as super_cursor:
+                    super_cursor.execute('SELECT current_user')
+                    self.assertEqual('postgres', super_cursor.fetchone()[0])
+
+                cursor.execute('SELECT current_user')
+                self.assertEqual('travis', cursor.fetchone()[0])
+
     def test_get_settings_from_entry_points(self):
         from ..utils import get_settings_from_entry_points
 

dbmigrator/utils.py

@@ -11,6 +11,7 @@
 except ImportError:
     # python 2
     import ConfigParser as configparser
+from contextlib import contextmanager
 import datetime
 import difflib
 import functools
@@ -36,6 +37,28 @@
 logger.addHandler(handler)
 
 
+_settings = {}
+
+
+def get_settings():
+    return _settings
+
+
+def set_settings(settings):
+    global _settings
+    _settings = settings
+
+
+@contextmanager
+def super_user():
+    settings = get_settings()
+    super_user = settings.get('super_user', 'postgres')
+    with psycopg2.connect(settings['db_connection_string'],
+                          user=super_user) as db_conn:
+        with db_conn.cursor() as cursor:
+            yield cursor
+
+
 # psycopg2 / libpq doesn't respond to SIGINT (ctrl-c):
 # https://github.com/psycopg/psycopg2/issues/333
 # To get around this problem, using code from: