Merge lp:~ricardokirkner/django-pgtools/fix-role-username-unescaped into lp:django-pgtools

[Ricardo Kirkner]

escape role and username properly

[Daniel Manrique (roadmr)]

Weird! LGTM, I tried to ensure the quote shuffling looks OK but ooh, all those quotes/double quotes...

review: Approve

[lp:~ricardokirkner/django-pgtools/fix-role-username-unescaped]

13. By Ricardo Kirkner on 2017-04-05

better quoting of identifiers

[William Grant (wgrant)]

review: Approve (code)

[lp:~ricardokirkner/django-pgtools/fix-role-username-unescaped]

14. By Ricardo Kirkner on 2017-04-05

make sure to respect psycopg2>=2.7 requirement in tox.ini

=== modified file 'pgtools/dbrole.py'
--- pgtools/dbrole.py	2017-01-16 19:08:45 +0000
+++ pgtools/dbrole.py	2017-04-05 13:23:18 +0000
@@ -5,6 +5,8 @@
 from django.core.management.base import CommandError
 from django.db import connection, transaction
 
+from psycopg2.extensions import quote_ident
+
 
 class DatabaseRole(object):
     ROLE_NAME = 'rolname'
@@ -23,7 +25,8 @@
                 raise CommandError("Role already exists: %s" % rolename)
 
             cursor = self.connection.cursor()
-            cursor.execute("CREATE ROLE %s" % rolename)
+            role = quote_ident(rolename, cursor.cursor)
+            cursor.execute('CREATE ROLE %s' % role)
             if commit and not self.exists(rolename):
                 raise CommandError("Role couldn't be created: %s" % rolename)
         elif not exists:
@@ -31,7 +34,8 @@
 
     def delete(self, commit=True):
         cursor = self.connection.cursor()
-        cursor.execute("DROP ROLE %s" % self.rolename)
+        role = quote_ident(self.rolename, cursor.cursor)
+        cursor.execute('DROP ROLE %s' % role)
         if commit and self.exists(self.rolename):
             raise CommandError("Role cannot be deleted: %s" % self.rolename)
 
@@ -40,19 +44,19 @@
         db_connection = connection if db_connection is None else db_connection
 
         cursor = db_connection.cursor()
-        cursor.execute('SELECT * FROM pg_roles WHERE rolname = %s', (rolename,))
+        cursor.execute('SELECT * FROM pg_roles WHERE rolname=%s', (rolename,))
         role = cursor.fetchone()
         return role is not None
 
     def get_users(self, order_by=ROLE_NAME):
         cursor = self.connection.cursor()
-        query = ("SELECT rolname FROM pg_authid "
-                 " WHERE oid IN ("
-                 "       SELECT member FROM pg_auth_members "
-                 "        WHERE roleid = ("
-                 "              SELECT oid FROM pg_authid "
-                 "               WHERE rolname = %%s)) "
-                 "ORDER BY %s" % order_by)
+        query = ('SELECT rolname FROM pg_authid '
+                 ' WHERE oid IN ('
+                 '       SELECT member FROM pg_auth_members '
+                 '        WHERE roleid = ('
+                 '              SELECT oid FROM pg_authid '
+                 '               WHERE rolname=%%s)) '
+                 'ORDER BY %s' % order_by)
         cursor.execute(query, (self.rolename,))
         users_result = cursor.fetchall()
         users = [result[0] for result in users_result]
@@ -60,10 +64,14 @@
 
     def grant(self, user, commit=True):
         cursor = self.connection.cursor()
-        cursor.execute("GRANT %s TO %s" % (self.rolename, user.username))
+        role = quote_ident(self.rolename, cursor.cursor)
+        username = quote_ident(user.username, cursor.cursor)
+        cursor.execute('GRANT %s TO %s' % (role, username))
 
     def revoke(self, user, commit=True):
         cursor = self.connection.cursor()
-        cursor.execute("REVOKE %s FROM %s" % (self.rolename, user.username))
+        role = quote_ident(self.rolename, cursor.cursor)
+        username = quote_ident(user.username, cursor.cursor)
+        cursor.execute('REVOKE %s FROM %s' % (role, username))
         if commit and user.has_role(self):
             raise CommandError("Role cannot be revoked: %s" % self.rolename)

=== modified file 'pgtools/dbuser.py'
--- pgtools/dbuser.py	2017-01-16 19:08:45 +0000
+++ pgtools/dbuser.py	2017-04-05 13:23:18 +0000
@@ -5,6 +5,8 @@
 from django.core.management.base import CommandError
 from django.db import connection, transaction
 
+from psycopg2.extensions import quote_ident
+
 
 class DatabaseUser(object):
 
@@ -22,11 +24,12 @@
                 raise CommandError("User already exists: %s" % username)
 
             cursor = self.connection.cursor()
+            user = quote_ident(username, cursor.cursor)
             if password is None:
-                cursor.execute("CREATE USER %s" % username)
+                cursor.execute('CREATE USER %s' % user)
             else:
-                cursor.execute("CREATE USER %s PASSWORD %%s" % username,
-                    (password,))
+                cursor.execute('CREATE USER %s PASSWORD %%s' % user,
+                               (password,))
             if commit and not self.exists(username):
                 raise CommandError(
                     "User couldn't be created: %s" % username)
@@ -35,7 +38,8 @@
 
     def delete(self, commit=True):
         cursor = self.connection.cursor()
-        cursor.execute("DROP USER %s" % self.username)
+        user = quote_ident(self.username, cursor.cursor)
+        cursor.execute('DROP USER %s' % user)
         if commit and self.exists(self.username):
             raise CommandError("User cannot be deleted: %s" % self.username)
 
@@ -43,7 +47,7 @@
     def exists(cls, username, db_connection=None):
         db_connection = connection if db_connection is None else db_connection
         cursor = db_connection.cursor()
-        cursor.execute('SELECT * FROM pg_roles WHERE rolname = %s', (username,))
+        cursor.execute('SELECT * FROM pg_roles WHERE rolname=%s', (username,))
         user = cursor.fetchone()
         return user is not None
 
@@ -52,10 +56,10 @@
         query = ('SELECT * FROM pg_auth_members '
                  ' WHERE roleid = ('
                  '       SELECT oid FROM pg_authid '
-                 '        WHERE rolname = %s) '
+                 '        WHERE rolname=%s) '
                  '   AND member = ('
                  '       SELECT oid FROM pg_authid '
-                 '        WHERE rolname = %s)')
+                 '        WHERE rolname=%s)')
         cursor.execute(query, (role.rolename, self.username))
         membership = cursor.fetchone()
         return bool(membership)

=== modified file 'pgtools/management/commands/grantuser.py'
--- pgtools/management/commands/grantuser.py	2017-04-03 13:41:34 +0000
+++ pgtools/management/commands/grantuser.py	2017-04-05 13:23:18 +0000
@@ -4,6 +4,7 @@
 
 from django.core.management.base import BaseCommand
 from django.db import connection
+from psycopg2.extensions import quote_ident
 
 from pgtools.decorators import graceful_db_errors
 from pgtools.utils import (
@@ -14,8 +15,8 @@
 )
 
 
-GRANT_SQL = "GRANT ALL ON %s TO %s;"
-EXISTS_SQL = "SELECT 1 FROM pg_class WHERE relname='%s';"
+GRANT_SQL = 'GRANT ALL ON %s TO %s'
+EXISTS_SQL = 'SELECT 1 FROM pg_class WHERE relname=%s'
 
 
 class Command(BaseCommand):
@@ -58,20 +59,20 @@
             connection.connection.commit()
 
     def _check_exists(self, table_or_sequence):
-        sql = EXISTS_SQL % table_or_sequence
-        self.cursor.execute(sql)
+        self.cursor.execute(EXISTS_SQL, (table_or_sequence,))
         result = self.cursor.fetchone()
         return bool(result)
 
     def _grant_one(self, table_or_sequence):
         print('Granting access to {}'.format(table_or_sequence))
-        sql = GRANT_SQL % (table_or_sequence, self.username)
-        self.cursor.execute(sql)
+        table = quote_ident(table_or_sequence, self.cursor.cursor)
+        username = quote_ident(self.username, self.cursor.cursor)
+        self.cursor.execute(GRANT_SQL % (table, username))
 
     def _grant_many_to_many(self, field):
         table = field.m2m_db_table()
-        self.cursor.execute("select pg_get_serial_sequence('%s', 'id')" %
-            table)
+        self.cursor.execute(
+            "SELECT pg_get_serial_sequence(%s, %s)", (table, 'id'))
         sequence = self.cursor.fetchone()[0]
         self._grant_one(table)
         self._grant_one(sequence)

=== modified file 'pgtools/tests.py'
--- pgtools/tests.py	2017-04-03 13:50:04 +0000
+++ pgtools/tests.py	2017-04-05 13:23:18 +0000
@@ -22,8 +22,8 @@
 from django.db.utils import load_backend
 from django.test import TransactionTestCase
 from django.test.utils import override_settings
-from mock import Mock, patch
-from psycopg2 import ProgrammingError
+from mock import Mock, call, patch
+from psycopg2 import ProgrammingError, errorcodes
 
 from pgtools.dbrole import DatabaseRole
 from pgtools.dbuser import DatabaseUser
@@ -38,10 +38,10 @@
     parse_username_and_rolename,
 )
 
-EMPTY_ROLE_NAME = 'test_empty_role'
-MASTER_ROLE_NAME = 'test_master_role'
-TEST_USER_NAME1 = 'test_user1'
-TEST_USER_NAME2 = 'test_user2'
+EMPTY_ROLE_NAME = 'test-empty_role'
+MASTER_ROLE_NAME = 'test-master_role'
+TEST_USER_NAME1 = 'test_user-1'
+TEST_USER_NAME2 = 'test_user-2'
 
 
 class CommandBaseTestCase(TransactionTestCase):
@@ -309,9 +309,9 @@
     def test_graceful_db_errors_permission_denied(self):
         @graceful_db_errors
         def test_func():
-            err = ProgrammingError()
-            err.pgcode = '42501'
-            raise err
+            class Error(ProgrammingError):
+                pgcode = errorcodes.INSUFFICIENT_PRIVILEGE
+            raise Error
         try:
             test_func()
             self.fail('No errors raised.')
@@ -322,9 +322,9 @@
     def test_graceful_db_errors_other(self):
         @graceful_db_errors
         def test_func():
-            err = ProgrammingError()
-            err.pgcode = 100
-            raise err
+            class Error(ProgrammingError):
+                pgcode = 100
+            raise Error
         try:
             test_func()
             self.fail('No errors raised.')
@@ -635,20 +635,42 @@
         self.assertEqual(len(mock_grant_one.call_args_list), 1)
 
     @patch('pgtools.management.commands.grantuser.get_models')
-    @patch('pgtools.management.commands.grantuser.connection')
-    def test__grant_many_to_many(self, mock_connection, mock_get_models):
-        mock_connection.cursor.return_value.fetchone.return_value = ['foo']
-        mock_get_models.return_value = self.get_models('auth')
-
-        call_command('grantuser', 'payments')
-
-        args_list = mock_connection.cursor.return_value.execute.call_args_list
-        calls = [x[0][0] for x in args_list]
-        for expected in [
-                "select pg_get_serial_sequence('auth_user_groups', 'id')",
-                'GRANT ALL ON foo TO payments;'
-            ]:
-            self.assertTrue(expected in calls)
+    def test__grant_many_to_many(self, mock_get_models):
+        sql = ["CREATE TABLE pgtools_bar (id integer PRIMARY KEY, baz integer);",
+               "CREATE TABLE pgtools_baz (id integer PRIMARY KEY, foo integer);"]
+        cursor = connection.cursor()
+        for query in sql:
+            cursor.execute(query)
+
+        class Bar(models.Model):
+            baz = models.ManyToManyField('Baz')
+
+        class Baz(models.Model):
+            foo = models.IntegerField()
+
+        mock_get_models.return_value = [Bar, Baz]
+
+        # make sure role exists
+        DatabaseRole('payments')
+
+        name = 'pgtools.management.commands.grantuser.connection.cursor'
+        with patch(name) as mock_cursor:
+            mock_cursor.return_value.cursor = cursor.cursor
+            mock_cursor.return_value.fetchone.return_value = [
+                'pgtools_bar_baz_id_seq'
+            ]
+            call_command('grantuser', 'payments')
+
+        mock_execute = mock_cursor.return_value.execute
+        expected_calls = [
+            call('SELECT pg_get_serial_sequence(%s, %s)',
+                 ('pgtools_bar_baz', 'id')),
+            call('GRANT ALL ON "pgtools_bar_baz" TO "payments"'),
+            call('GRANT ALL ON "pgtools_bar_baz_id_seq" TO "payments"'),
+        ]
+
+        for call_ in expected_calls:
+            self.assertIn(call_, mock_execute.call_args_list)
 
 
 @skipIf(DJANGO_VERSION > (1, 6), "Django 1.7 and later use django.apps for getting models")

=== modified file 'requirements.txt'
--- requirements.txt	2017-01-16 19:01:08 +0000
+++ requirements.txt	2017-04-05 13:23:18 +0000
@@ -1,3 +1,3 @@
 mock
-psycopg2==2.4.5
+psycopg2>=2.7
 django

=== modified file 'tox.ini'
--- tox.ini	2017-04-03 13:41:34 +0000
+++ tox.ini	2017-04-05 13:23:18 +0000
@@ -6,7 +6,7 @@
 basepython =
     python2.7
 deps =
-    psycopg2 == 2.4.5
+    psycopg2 >= 2.7
     mock
 
 [testenv:py3]