Add grant_priviliege method to the client.

geodimm · geodimm · commit 36e50eacf3e3 · 2015-05-04T20:42:07.000+01:00
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -642,6 +642,22 @@ def revoke_admin_privileges(self, username):
         text = "REVOKE ALL PRIVILEGES FROM {}".format(username)
         self.query(text)
 
+    def grant_privilege(self, privilege, database, username):
+        """Grant a privilege on a database to an user.
+
+        :param privilege: the privilege to grant, one of 'read', 'write'
+            or 'all'. The string is case-insensitive
+        :type privilege: str
+        :param database: the database to grant the privilege on
+        :type database: str
+        :param username: the username to grant the privilege to
+        :type username: str
+        """
+        text = "GRANT {} ON {} TO {}".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
@@ -668,6 +668,28 @@ def test_revoke_admin_privileges_invalid(self):
         with _mocked_session(cli, 'get', 400):
             self.cli.revoke_admin_privileges('')
 
+    def test_grant_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.grant_privilege('read', 'testdb', 'test')
+
+            self.assertEqual(
+                m.last_request.qs['q'][0],
+                'grant read on testdb to test'
+            )
+
+    @raises(Exception)
+    def test_grant_privilege_invalid(self):
+        cli = InfluxDBClient('host', 8086, 'username', 'password')
+        with _mocked_session(cli, 'get', 400):
+            self.cli.grant_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
@@ -407,6 +407,21 @@ def test_revoke_admin_privileges_invalid(self):
         self.assertIn('{"error":"error parsing query: ',
                       ctx.exception.content)
 
+    def test_grant_privilege(self):
+        self.cli.create_user('test', 'test')
+        self.cli.create_database('testdb')
+        self.cli.grant_privilege('all', 'testdb', 'test')
+        # TODO: when supported by InfluxDB, check if privileges are granted
+
+    def test_grant_privilege_invalid(self):
+        self.cli.create_user('test', 'test')
+        self.cli.create_database('testdb')
+        with self.assertRaises(InfluxDBClientError) as ctx:
+            self.cli.grant_privilege('', 'testdb', 'test')
+        self.assertEqual(400, ctx.exception.code)
+        self.assertIn('{"error":"error parsing query: ',
+                      ctx.exception.content)
+
 
 ############################################################################
 
