Merge lp:~ricardokirkner/django-pgtools/fix-grantuser-sequence-fully-qualified into lp:django-pgtools

[Ricardo Kirkner]

fix granting on sequence from m2m model when pg returns the fully qualified name

[Ricardo Kirkner (ricardokirkner)]

review: Approve

=== modified file 'pgtools/management/commands/grantuser.py'
--- pgtools/management/commands/grantuser.py	2017-04-05 13:14:23 +0000
+++ pgtools/management/commands/grantuser.py	2017-04-05 18:30:10 +0000
@@ -73,7 +73,13 @@
         table = field.m2m_db_table()
         self.cursor.execute(
             "SELECT pg_get_serial_sequence(%s, %s)", (table, 'id'))
-        sequence = self.cursor.fetchone()[0]
+        # if the sequence name is returned as a fully qualified name
+        # ie, like <schema>.<name>, only return the name part
+        result = self.cursor.fetchone()[0]
+        if '.' in result:
+            sequence = result.split('.')[1]
+        else:
+            sequence = result
         self._grant_one(table)
         self._grant_one(sequence)
 
 
=== modified file 'pgtools/tests.py'
--- pgtools/tests.py	2017-04-05 13:14:23 +0000
+++ pgtools/tests.py	2017-04-05 18:30:10 +0000
@@ -636,12 +636,6 @@
 
     @patch('pgtools.management.commands.grantuser.get_models')
     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')
 
@@ -651,8 +645,10 @@
         mock_get_models.return_value = [Bar, Baz]
 
         # make sure role exists
-        DatabaseRole('payments')
+        if not DatabaseRole.exists('payments'):
+            DatabaseRole('payments', create=True)
 
+        cursor = connection.cursor()
         name = 'pgtools.management.commands.grantuser.connection.cursor'
         with patch(name) as mock_cursor:
             mock_cursor.return_value.cursor = cursor.cursor
@@ -672,6 +668,41 @@
         for call_ in expected_calls:
             self.assertIn(call_, mock_execute.call_args_list)
 
+    @patch('pgtools.management.commands.grantuser.get_models')
+    def test__grant_many_to_many_with_sequence_full_name(
+            self, mock_get_models):
+        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
+        if not DatabaseRole.exists('payments'):
+            DatabaseRole('payments', create=True)
+
+        cursor = connection.cursor()
+        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 = [
+                'public.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")
 class Django16GrantUserCommandTestCase(GrantUserCommandTestCase):