simplify design point fetch queries

1 files changed, 5 insertions, 7 deletions

diff --git a/rtss14/database.py b/rtss14/database.py
index 6028e3a..0f9a917 100755
--- a/rtss14/database.py
+++ b/rtss14/database.py
@@ -351,22 +351,20 @@ def __get_design_points(conn, ndp):
     ####################
     try:
         begin_sync(conn, c)
-        c.execute('SELECT id FROM dp_pending WHERE taken=0 LIMIT %s FOR UPDATE', (ndp,))
-        dp_ids = [d['id'] for d in c.fetchall()]
-        nfetched = len(dp_ids) if dp_ids else 0
+        c.execute('SELECT id,%s FROM dp_pending WHERE taken=0 LIMIT %%s FOR UPDATE' % ','.join(dp_col_names), (ndp,))
+        rows = c.fetchall()
+        nfetched = c.rowcount
         if nfetched > 0:
+            dp_ids = [r['id'] for r in rows]
             c.execute('UPDATE dp_pending SET taken=1 WHERE id IN (%s)' % ','.join(['%s']*len(dp_ids)), tuple(dp_ids))
         end_sync(conn)
+        fetched = [dp_to_dp(r) for r in rows]
     except db.OperationalError, e:
         print e
         print 'Last Query: %s' % c._last_executed
         rollback(conn)
         c.close()
         raise e
-    ####################
-    if nfetched > 0:
-        c.execute('SELECT %s FROM dp_pending WHERE id IN (%s)' % (','.join(dp_col_names), ','.join(['%s']*len(dp_ids))), tuple(dp_ids))
-        fetched = [db_to_dp(d) for d in c.fetchall()]
     c.close()
     return fetched