Add a multi-column query wrapper _get_element2()

Obviously modelled on _get_element()

Change-Id: I5dd5f5b2fa8ac96bd9c8038d41eb13f45f38d85b

1 files changed, 24 insertions, 0 deletions

diff --git a/db/db.py b/db/db.py
index 108b2a2..0708cb6 100644
--- a/db/db.py
+++ b/db/db.py
@@ -213,6 +213,30 @@ class VlanDB:
         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):
+
+        # We really want to use psycopg's type handling deal with the
+        # (potentially) user-supplied data in the value field, so we
+        # have to pass (sql,data) through to cursor.execute. However,
+        # we can't have psycopg do all the argument substitution here
+        # 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")
+
+        # 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)
+        self.cursor.execute(sql, data)
+
+        if self.cursor.rowcount > 0:
+            return self.cursor.fetchone()[0]
+        else:
+            return None
+
                 results.append(record[0])
         else:
             return None