Add a new wrapper _get_multi_elements()

Return a dict of multiple entries from a query on one column

Change-Id: I3f639220a21758d0e870be3236ff9f8d1d20602b

1 files changed, 15 insertions, 14 deletions

diff --git a/db/db.py b/db/db.py
index c360be1..870619c 100644
--- a/db/db.py
+++ b/db/db.py
@@ -196,8 +196,8 @@ class VlanDB:
         else:
             return None
 
-    # Grab one column from multiple rows of a query; useful as a quick wrapper
-    def _get_multi_elements(self, select_field, table, compare_field, value):
+    # Grab one column from one row of a query on 2 columns; useful as a quick wrapper
+    def _get_element2(self, select_field, table, compare_field1, value1, compare_field2, value2):
 
         # We really want to use psycopg's type handling deal with the
         # (potentially) user-supplied data in the value field, so we
@@ -206,20 +206,22 @@ class VlanDB:
         # as it will quote all the params like the table name. That
         # doesn't work. So, we substitute a "%s" for "%s" here so we
         # keep it after python's own string substitution.
-        sql = "SELECT %s FROM %s WHERE %s = %s" % (select_field, table, compare_field, "%s")
+        sql = "SELECT %s FROM %s WHERE %s = %s AND %s = %s" % (select_field, table, compare_field1, "%s", compare_field2, "%s")
 
         # Now, the next icky thing: we need to make sure that we're
         # passing a dict so that psycopg2 can pick it apart properly
         # for its own substitution code. We force this with the
         # trailing comma here
-        data = (value, )
+        data = (value1, value2)
         self.cursor.execute(sql, data)
 
         if self.cursor.rowcount > 0:
-            results = []
-            for record in self.cursor:
-    # Grab one column from one row of a query on 2 columns; useful as a quick wrapper
-    def _get_element2(self, select_field, table, compare_field1, value1, compare_field2, value2):
+            return self.cursor.fetchone()[0]
+        else:
+            return None
+
+    # Grab one column from multiple rows of a query; useful as a quick wrapper
+    def _get_multi_elements(self, select_field, table, compare_field, value):
 
         # We really want to use psycopg's type handling deal with the
         # (potentially) user-supplied data in the value field, so we
@@ -228,21 +230,20 @@ class VlanDB:
         # as it will quote all the params like the table name. That
         # doesn't work. So, we substitute a "%s" for "%s" here so we
         # keep it after python's own string substitution.
-        sql = "SELECT %s FROM %s WHERE %s = %s AND %s = %s" % (select_field, table, compare_field1, "%s", compare_field2, "%s")
+        sql = "SELECT %s FROM %s WHERE %s = %s" % (select_field, table, compare_field, "%s")
 
         # Now, the next icky thing: we need to make sure that we're
         # passing a dict so that psycopg2 can pick it apart properly
         # for its own substitution code. We force this with the
         # trailing comma here
-        data = (value1, value2)
+        data = (value, )
         self.cursor.execute(sql, data)
 
         if self.cursor.rowcount > 0:
-            return self.cursor.fetchone()[0]
-        else:
-            return None
-
+            results = []
+            for record in self.cursor:
                 results.append(record[0])
+            return results
         else:
             return None