Reduce boilerplate in db functions with context managers

youtube/subscriptions.py

@@ -8,6 +8,7 @@ import gevent
 import html
 import json
 import traceback
+import contextlib
 
 with open('yt_subscriptions_template.html', 'r', encoding='utf-8') as f:
     subscriptions_template = Template(f.read())
@@ -58,7 +59,8 @@ def open_database():
         connection.close()
         raise
 
-    return connection
+    # https://stackoverflow.com/questions/19522505/using-sqlite3-in-python-with-with-keyword
+    return contextlib.closing(connection)
 
 def _subscribe(channels):
     ''' channels is a list of (channel_id, channel_name) '''
@@ -66,35 +68,21 @@ def _subscribe(channels):
     # set time_last_checked to 0 on all channels being subscribed to
     channels = ( (channel_id, channel_name, 0) for channel_id, channel_name in channels)
 
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
+            cursor.executemany('''INSERT OR IGNORE INTO subscribed_channels (yt_channel_id, channel_name, time_last_checked)
-        cursor.executemany("INSERT OR IGNORE INTO subscribed_channels (yt_channel_id, channel_name, time_last_checked) VALUES (?, ?, ?)", channels)
+                                  VALUES (?, ?, ?)''', channels)
-        connection.commit()
-    except:
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
 
 # TODO: delete thumbnails
 def _unsubscribe(channel_ids):
     ''' channel_ids is a list of channel_ids '''
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
             cursor.executemany("DELETE FROM subscribed_channels WHERE yt_channel_id=?", ((channel_id, ) for channel_id in channel_ids))
-        connection.commit()
-    except:
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
 
 def _get_videos(number, offset):
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
             db_videos = cursor.execute('''SELECT video_id, title, duration, channel_name
                                           FROM videos
                                           INNER JOIN subscribed_channels on videos.sql_channel_id = subscribed_channels.id
@@ -108,104 +96,60 @@ def _get_videos(number, offset):
                     'duration': db_video[2],
                     'author':   db_video[3],
                 }
-    except:
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
 
 def _get_subscribed_channels():
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
             for item in cursor.execute('''SELECT channel_name, yt_channel_id
                                           FROM subscribed_channels
                                           ORDER BY channel_name'''):
                 yield item
-    except:
+
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
 
 def _add_tags(channel_ids, tags):
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
-
             pairs = [(tag, yt_channel_id) for tag in tags for yt_channel_id in channel_ids]
-
             cursor.executemany('''INSERT OR IGNORE INTO tag_associations (tag, sql_channel_id)
                                   SELECT ?, id FROM subscribed_channels WHERE yt_channel_id = ? ''', pairs)
-        connection.commit()
-    except:
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
-
 
 
 def _remove_tags(channel_ids, tags):
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
-
             pairs = [(tag, yt_channel_id) for tag in tags for yt_channel_id in channel_ids]
-
+            cursor.executemany('''DELETE FROM tag_associations
-        cursor.executemany('''DELETE FROM tag_associations WHERE tag = ? AND sql_channel_id = (
+                                  WHERE tag = ? AND sql_channel_id = (
                                       SELECT id FROM subscribed_channels WHERE yt_channel_id = ?
                                    )''', pairs)
-        connection.commit()
+
-    except:
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
 
 
 def _get_tags(channel_id):
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
-
             return [row[0] for row in cursor.execute('''SELECT tag
                                                         FROM tag_associations
                                                         WHERE sql_channel_id = (
                                                             SELECT id FROM subscribed_channels WHERE yt_channel_id = ?
                                                         )''', (channel_id,))]
-    except:
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
 
 def _get_all_tags():
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
             return [row[0] for row in cursor.execute('''SELECT DISTINCT tag FROM tag_associations''')]
-    except:
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
 
 def _get_channel_names(channel_ids):
     ''' returns list of (channel_id, channel_name) '''
-    connection = open_database()
+    with open_database() as connection:
-    try:
+        with connection as cursor:
-        cursor = connection.cursor()
             result = []
             for channel_id in channel_ids:
-            row = cursor.execute('''SELECT channel_name FROM subscribed_channels WHERE yt_channel_id = ?''', (channel_id,)).fetchone()
+                row = cursor.execute('''SELECT channel_name
+                                        FROM subscribed_channels
+                                        WHERE yt_channel_id = ?''', (channel_id,)).fetchone()
                 result.append( (channel_id, row[0]) )
             return result
-    except:
-        connection.rollback()
-        raise
-    finally:
-        connection.close()
 
 
 units = {
@@ -464,20 +408,13 @@ def post_subscriptions_page(env, start_response):
         _unsubscribe(params['channel_id'])
 
     elif action == 'refresh':
-        connection = open_database()
+        with open_database() as connection:
-        try:
+            with connection as cursor:
-            cursor = connection.cursor()
                 for sql_channel_id, yt_channel_id in cursor.execute('''SELECT id, yt_channel_id FROM subscribed_channels''').fetchall():
                     db_videos = ( (sql_channel_id, info['id'], info['title'], info['duration'], info['time_published'], info['description']) for info in _get_upstream_videos(yt_channel_id) )
                     cursor.executemany('''INSERT OR IGNORE INTO videos (sql_channel_id, video_id, title, duration, time_published, description) VALUES (?, ?, ?, ?, ?, ?)''', db_videos)
 
                 cursor.execute('''UPDATE subscribed_channels SET time_last_checked = ?''', ( int(time.time()), ) )
-            connection.commit()
-        except:
-            connection.rollback()
-            raise
-        finally:
-            connection.close()
 
         start_response('303 See Other', [('Location', util.URL_ORIGIN + '/subscriptions'),] )
         return b''