Merge pull request influxdata#361 from saaros/quote_ident (Thanks @saaros!)

aviau · web-flow · commit 898c4b68e3f8 · 2016-09-09T11:30:35.000-04:00
Properly quote all identifiers and literals in InfluxDBClient
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -13,7 +13,7 @@
 import requests.exceptions
 from sys import version_info
 
-from influxdb.line_protocol import make_lines
+from influxdb.line_protocol import make_lines, quote_ident, quote_literal
 from influxdb.resultset import ResultSet
 from .exceptions import InfluxDBClientError
 from .exceptions import InfluxDBServerError
@@ -552,8 +552,8 @@ def alter_retention_policy(self, name, database=None,
             should be set. Otherwise the operation will fail.
         """
         query_string = (
-            "ALTER RETENTION POLICY \"{0}\" ON \"{1}\""
-        ).format(name, database or self._database)
+            "ALTER RETENTION POLICY {0} ON {1}"
+        ).format(quote_ident(name), quote_ident(database or self._database))
         if duration:
             query_string += " DURATION {0}".format(duration)
         if replication:
@@ -573,8 +573,8 @@ def drop_retention_policy(self, name, database=None):
         :type database: str
         """
         query_string = (
-            "DROP RETENTION POLICY \"{0}\" ON \"{1}\""
-        ).format(name, database or self._database)
+            "DROP RETENTION POLICY {0} ON {1}"
+        ).format(quote_ident(name), quote_ident(database or self._database))
         self.query(query_string)
 
     def get_list_retention_policies(self, database=None):
@@ -631,8 +631,8 @@ def create_user(self, username, password, admin=False):
             privileges or not
         :type admin: boolean
         """
-        text = "CREATE USER \"{0}\" WITH PASSWORD '{1}'".format(username,
-                                                                password)
+        text = "CREATE USER {0} WITH PASSWORD {1}".format(
+            quote_ident(username), quote_literal(password))
         if admin:
             text += ' WITH ALL PRIVILEGES'
         self.query(text)
@@ -643,7 +643,7 @@ def drop_user(self, username):
         :param username: the username to drop
         :type username: str
         """
-        text = "DROP USER {0}".format(username)
+        text = "DROP USER {0}".format(quote_ident(username))
         self.query(text)
 
     def set_user_password(self, username, password):
@@ -654,7 +654,8 @@ def set_user_password(self, username, password):
         :param password: the new password for the user
         :type password: str
         """
-        text = "SET PASSWORD FOR {0} = '{1}'".format(username, password)
+        text = "SET PASSWORD FOR {0} = {1}".format(
+            quote_ident(username), quote_literal(password))
         self.query(text)
 
     def delete_series(self, database=None, measurement=None, tags=None):
@@ -672,11 +673,12 @@ def delete_series(self, database=None, measurement=None, tags=None):
         database = database or self._database
         query_str = 'DROP SERIES'
         if measurement:
-            query_str += ' FROM "{0}"'.format(measurement)
+            query_str += ' FROM {0}'.format(quote_ident(measurement))
 
         if tags:
-            query_str += ' WHERE ' + ' and '.join(["{0}='{1}'".format(k, v)
-                                                   for k, v in tags.items()])
+            tag_eq_list = ["{0}={1}".format(quote_ident(k), quote_literal(v))
+                           for k, v in tags.items()]
+            query_str += ' WHERE ' + ' AND '.join(tag_eq_list)
         self.query(query_str, database=database)
 
     def grant_admin_privileges(self, username):
@@ -688,7 +690,7 @@ def grant_admin_privileges(self, username):
         .. note:: Only a cluster administrator can create/drop databases
             and manage users.
         """
-        text = "GRANT ALL PRIVILEGES TO {0}".format(username)
+        text = "GRANT ALL PRIVILEGES TO {0}".format(quote_ident(username))
         self.query(text)
 
     def revoke_admin_privileges(self, username):
@@ -700,7 +702,7 @@ def revoke_admin_privileges(self, username):
         .. note:: Only a cluster administrator can create/ drop databases
             and manage users.
         """
-        text = "REVOKE ALL PRIVILEGES FROM {0}".format(username)
+        text = "REVOKE ALL PRIVILEGES FROM {0}".format(quote_ident(username))
         self.query(text)
 
     def grant_privilege(self, privilege, database, username):
@@ -715,8 +717,8 @@ def grant_privilege(self, privilege, database, username):
         :type username: str
         """
         text = "GRANT {0} ON {1} TO {2}".format(privilege,
-                                                database,
-                                                username)
+                                                quote_ident(database),
+                                                quote_ident(username))
         self.query(text)
 
     def revoke_privilege(self, privilege, database, username):
@@ -731,8 +733,8 @@ def revoke_privilege(self, privilege, database, username):
         :type username: str
         """
         text = "REVOKE {0} ON {1} FROM {2}".format(privilege,
-                                                   database,
-                                                   username)
+                                                   quote_ident(database),
+                                                   quote_ident(username))
         self.query(text)
 
     def get_list_privileges(self, username):
@@ -754,7 +756,7 @@ def get_list_privileges(self, username):
              {u'privilege': u'ALL PRIVILEGES', u'database': u'db2'},
              {u'privilege': u'NO PRIVILEGES', u'database': u'db3'}]
         """
-        text = "SHOW GRANTS FOR {0}".format(username)
+        text = "SHOW GRANTS FOR {0}".format(quote_ident(username))
         return list(self.query(text).get_points())
 
     def send_packet(self, packet, protocol='json'):
diff --git a/influxdb/line_protocol.py b/influxdb/line_protocol.py
@@ -53,16 +53,32 @@ def _escape_tag(tag):
     )
 
 
+def quote_ident(value):
+    return "\"{0}\"".format(
+        value.replace(
+            "\\", "\\\\"
+        ).replace(
+            "\"", "\\\""
+        ).replace(
+            "\n", "\\n"
+        )
+    )
+
+
+def quote_literal(value):
+    return "'{0}'".format(
+        value.replace(
+            "\\", "\\\\"
+        ).replace(
+            "'", "\\'"
+        )
+    )
+
+
 def _escape_value(value):
     value = _get_unicode(value)
     if isinstance(value, text_type) and value != '':
-        return "\"{0}\"".format(
-            value.replace(
-                "\"", "\\\""
-            ).replace(
-                "\n", "\\n"
-            )
-        )
+        return quote_ident(value)
     elif isinstance(value, integer_types) and not isinstance(value, bool):
         return str(value) + 'i'
     else:
diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py
@@ -688,7 +688,7 @@ def test_grant_admin_privileges(self):
 
             self.assertEqual(
                 m.last_request.qs['q'][0],
-                'grant all privileges to test'
+                'grant all privileges to "test"'
             )
 
     @raises(Exception)
@@ -710,7 +710,7 @@ def test_revoke_admin_privileges(self):
 
             self.assertEqual(
                 m.last_request.qs['q'][0],
-                'revoke all privileges from test'
+                'revoke all privileges from "test"'
             )
 
     @raises(Exception)
@@ -732,7 +732,7 @@ def test_grant_privilege(self):
 
             self.assertEqual(
                 m.last_request.qs['q'][0],
-                'grant read on testdb to test'
+                'grant read on "testdb" to "test"'
             )
 
     @raises(Exception)
@@ -754,7 +754,7 @@ def test_revoke_privilege(self):
 
             self.assertEqual(
                 m.last_request.qs['q'][0],
-                'revoke read on testdb from test'
+                'revoke read on "testdb" from "test"'
             )
 
     @raises(Exception)
diff --git a/influxdb/tests/server_tests/client_test_with_server.py b/influxdb/tests/server_tests/client_test_with_server.py
@@ -200,14 +200,6 @@ def test_drop_user_nonexisting(self):
         self.assertIn('user not found',
                       ctx.exception.content)
 
-    def test_drop_user_invalid(self):
-        with self.assertRaises(InfluxDBClientError) as ctx:
-            self.cli.drop_user('very invalid')
-        self.assertEqual(400, ctx.exception.code)
-        self.assertIn('{"error":"error parsing query: '
-                      'found invalid, expected',
-                      ctx.exception.content)
-
     @unittest.skip("Broken as of 0.9.0")
     def test_revoke_admin_privileges(self):
         self.cli.create_user('test', 'test', admin=True)
@@ -217,13 +209,6 @@ def test_revoke_admin_privileges(self):
         self.assertEqual([{'user': 'test', 'admin': False}],
                          self.cli.get_list_users())
 
-    def test_revoke_admin_privileges_invalid(self):
-        with self.assertRaises(InfluxDBClientError) as ctx:
-            self.cli.revoke_admin_privileges('')
-        self.assertEqual(400, ctx.exception.code)
-        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')
diff --git a/influxdb/tests/test_line_protocol.py b/influxdb/tests/test_line_protocol.py
@@ -76,3 +76,15 @@ def test_make_lines_unicode(self):
             line_protocol.make_lines(data),
             'test,unicode_tag=\'Привет!\' unicode_val="Привет!"\n'
         )
+
+    def test_quote_ident(self):
+        self.assertEqual(
+            line_protocol.quote_ident(r"""\foo ' bar " Örf"""),
+            r'''"\\foo ' bar \" Örf"'''
+        )
+
+    def test_quote_literal(self):
+        self.assertEqual(
+            line_protocol.quote_literal(r"""\foo ' bar " Örf"""),
+            r"""'\\foo \' bar " Örf'"""
+        )
