Add revoke_privilege method to the client.

geodimm · geodimm · commit 731c4e073011 · 2015-05-04T20:42:28.000+01:00
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -658,6 +658,22 @@ def grant_privilege(self, privilege, database, username):
                                              username)
         self.query(text)
 
+    def revoke_privilege(self, privilege, database, username):
+        """Revoke a privilege on a database from an user.
+
+        :param privilege: the privilege to revoke, one of 'read', 'write'
+            or 'all'. The string is case-insensitive
+        :type privilege: str
+        :param database: the database to revoke the privilege on
+        :type database: str
+        :param username: the username to revoke the privilege from
+        :type username: str
+        """
+        text = "REVOKE {} ON {} FROM {}".format(privilege,
+                                                database,
+                                                username)
+        self.query(text)
+
     def send_packet(self, packet):
         """Send an UDP packet.
 
diff --git a/tests/influxdb/client_test.py b/tests/influxdb/client_test.py
@@ -690,6 +690,28 @@ def test_grant_privilege_invalid(self):
         with _mocked_session(cli, 'get', 400):
             self.cli.grant_privilege('', 'testdb', 'test')
 
+    def test_revoke_privilege(self):
+        example_response = '{"results":[{}]}'
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(
+                requests_mock.GET,
+                "http://localhost:8086/query",
+                text=example_response
+            )
+            self.cli.revoke_privilege('read', 'testdb', 'test')
+
+            self.assertEqual(
+                m.last_request.qs['q'][0],
+                'revoke read on testdb from test'
+            )
+
+    @raises(Exception)
+    def test_revoke_privilege_invalid(self):
+        cli = InfluxDBClient('host', 8086, 'username', 'password')
+        with _mocked_session(cli, 'get', 400):
+            self.cli.revoke_privilege('', 'testdb', 'test')
+
 
 class FakeClient(InfluxDBClient):
     fail = False
diff --git a/tests/influxdb/client_test_with_server.py b/tests/influxdb/client_test_with_server.py
@@ -422,6 +422,21 @@ def test_grant_privilege_invalid(self):
         self.assertIn('{"error":"error parsing query: ',
                       ctx.exception.content)
 
+    def test_revoke_privilege(self):
+        self.cli.create_user('test', 'test')
+        self.cli.create_database('testdb')
+        self.cli.revoke_privilege('all', 'testdb', 'test')
+        # TODO: when supported by InfluxDB, check if privileges are revoked
+
+    def test_revoke_privilege_invalid(self):
+        self.cli.create_user('test', 'test')
+        self.cli.create_database('testdb')
+        with self.assertRaises(InfluxDBClientError) as ctx:
+            self.cli.revoke_privilege('', 'testdb', 'test')
+        self.assertEqual(400, ctx.exception.code)
+        self.assertIn('{"error":"error parsing query: ',
+                      ctx.exception.content)
+
 
 ############################################################################
 
