Merge lp:~ricardokirkner/django-pgtools/check-exists-before-grant into lp:django-pgtools

[Ricardo Kirkner]

Before requesting a GRANT operation, make sure the relation actually exists, to avoid causing errors on the cursor.

Before requesting a GRANT operation, make sure the relation actually exists, to avoid causing errors on the cursor.

This aims specifically at fixing the issue as seen on https://pastebin.canonical.com/85297/

[lp:~ricardokirkner/django-pgtools/check-exists-before-grant]

8. By Ricardo Kirkner on 2013-02-21

fixed review comments

- styling issue
- added test to verify result without mocking cursor.fetchone

[Matias Bordese (matiasb)]

LGTM

review: Approve

=== modified file 'pgtools/management/commands/grantuser.py'
--- pgtools/management/commands/grantuser.py	2012-04-13 17:32:45 +0000
+++ pgtools/management/commands/grantuser.py	2013-02-21 21:19:22 +0000
@@ -11,6 +11,7 @@
 
 
 GRANT_SQL = "GRANT ALL ON %s TO %s;"
+EXISTS_SQL = "SELECT 1 FROM pg_class WHERE relname='%s';"
 
 
 class Command(BaseCommand):
@@ -34,10 +35,13 @@
         for model in get_models():
             # grant access to model
             table = model._meta.db_table
-            self._grant_one(table)
-            if model._meta.has_auto_field:
-                sequence = "%s_%s_seq" % (table, model._meta.auto_field.column)
-                self._grant_one(sequence)
+            if self._check_exists(table):
+                self._grant_one(table)
+                if model._meta.has_auto_field:
+                    sequence = "%s_%s_seq" % (table,
+                                              model._meta.auto_field.column)
+                    if self._check_exists(sequence):
+                        self._grant_one(sequence)
 
             # grant access to model relations
             for f in model._meta.many_to_many:
@@ -46,6 +50,12 @@
         if not getattr(connection.features, 'uses_autocommit', False):
             connection.connection.commit()
 
+    def _check_exists(self, table_or_sequence):
+        sql = EXISTS_SQL % table_or_sequence
+        self.cursor.execute(sql)
+        result = self.cursor.fetchone()
+        return bool(result)
+
     def _grant_one(self, table_or_sequence):
         print 'Granting access to', table_or_sequence
         sql = GRANT_SQL % (table_or_sequence, self.username)

=== modified file 'pgtools/tests.py'
--- pgtools/tests.py	2013-02-21 15:01:28 +0000
+++ pgtools/tests.py	2013-02-21 21:19:22 +0000
@@ -17,7 +17,7 @@
 from django.conf import settings
 from django.core.management import call_command
 from django.core.management.base import CommandError
-from django.db import connection
+from django.db import connection, models
 from django.db.backends import postgresql_psycopg2
 from django.db.models.loading import get_models
 from django.db.utils import load_backend
@@ -570,6 +570,31 @@
             (('django_admin_log_id_seq',), {}))
 
     @patch('pgtools.management.commands.grantuser.get_models')
+    def test_command_grantuser_missing_relations(self, mock_get_models):
+        # create table without sequence
+        sql = "CREATE TABLE pgtools_foo (id integer PRIMARY KEY);"
+        cursor = connection.cursor()
+        cursor.execute(sql)
+
+        class Foo(models.Model):
+            foo = models.IntegerField()
+
+        mock_get_models.return_value = [Foo]
+        command = grantuser.Command
+        with patch.object(command, '_grant_one') as mock_grant_one:
+            # force the cache to reload the apps
+            try:
+                call_command('grantuser', 'payments')
+            except:
+                self.fail('Error while calling grantuser command')
+
+        # assert the grant table method was called on all model tables
+        self.assertEqual(mock_grant_one.call_args_list[0],
+            (('pgtools_foo',), {}))
+        # assert nothing else was processed
+        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']