Merge pull request #1200 from josenavas/queue-issue

Using regex rather than directly check for indices in the sql args

Author: squirrelo

qiita_db/sql_connection.py

@@ -81,6 +81,7 @@
 from itertools import chain
 from tempfile import mktemp
 from datetime import date, time, datetime
+import re
 
 from psycopg2 import (connect, ProgrammingError, Error as PostgresError,
                       OperationalError)
@@ -154,6 +155,8 @@ class SQLConnectionHandler(object):
                  'admin_with_database': '_admin_args',
                  'admin_without_database': '_admin_nodb_args'}
 
+    _regex = re.compile("{(\d+)}")
+
     """Encapsulates the DB connection with the Postgres DB
 
     Parameters
@@ -362,11 +365,12 @@ def execute_queue(self, queue):
                 if sql_args is not None:
                     for pos, arg in enumerate(sql_args):
                         # check if previous results needed and replace
-                        if isinstance(arg, str) and \
-                                arg[0] == "{" and arg[-1] == "}":
-                            result_pos = int(arg[1:-1])
-                            sql_args[pos] = results[result_pos]
-                            clear_res = True
+                        if isinstance(arg, str):
+                            result = self._regex.search(arg)
+                            if result:
+                                result_pos = int(result.group(1))
+                                sql_args[pos] = results[result_pos]
+                                clear_res = True
                 # wipe out results if needed and reset clear_res
                 if clear_res:
                     results = []

qiita_db/test/test_sql_connection.py

@@ -93,6 +93,25 @@ def test_run_queue_placeholders(self):
                 None, None, None]]
         self.assertEqual(obs, exp)
 
+    def test_run_queue_placeholders_regex(self):
+        self.conn_handler.create_queue("toy_queue")
+        self.conn_handler.add_to_queue(
+            "toy_queue", "INSERT INTO qiita.qiita_user (email, name, password,"
+            "phone) VALUES (%s, %s, %s, %s) RETURNING email",
+            ['insert@foo.bar', '', 'pass', '111-111-11112'])
+        self.conn_handler.add_to_queue(
+            "toy_queue", "UPDATE qiita.qiita_user SET user_level_id = 1, "
+            "phone = '222-222-2221' WHERE email = %s AND name = %s",
+            ['{0}', ''])
+        obs = self.conn_handler.execute_queue("toy_queue")
+        self.assertEqual(obs, [])
+        obs = self.conn_handler.execute_fetchall(
+            "SELECT * from qiita.qiita_user WHERE email = %s",
+            ['insert@foo.bar'])
+        exp = [['insert@foo.bar', 1, 'pass', '', None, None, '222-222-2221',
+                None, None, None]]
+        self.assertEqual(obs, exp)
+
     def test_queue_fail(self):
         """Fail if no results data exists for substitution"""
         self.conn_handler = SQLConnectionHandler()