Merge pull request #4 from ilri/database-refactor

Database refactor

CHANGELOG.md

@@ -10,6 +10,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Hound CI configuration to validate pull requests against PEP 8 code style with Flake8
 - Configuration for [pipenv](https://pipenv.readthedocs.io/en/latest/)
 
+## Changed
+- Use a database management class with Python context management to automatically open/close connections and cursors
+
 ## Changed
 - Validate code against PEP 8 style guide with Flake8
 

README.md

@@ -70,7 +70,6 @@ The item id is the *internal* id for an item. You can get these from the standar
 
 ## Todo
 
-- Close DB connection when gunicorn shuts down gracefully
 - Better logging
 - Tests
 - Check if database exists (try/except)

dspace_statistics_api/app.py

@@ -1,9 +1,6 @@
-from .database import database_connection
+from .database import DatabaseManager
 import falcon
 
-db = database_connection()
-db.set_session(readonly=True)
-
 
 class RootResource:
     def on_get(self, req, resp):
@@ -21,8 +18,10 @@ class AllItemsResource:
         page = req.get_param_as_int("page", min=0) or 0
         offset = limit * page
 
-        cursor = db.cursor()
+        with DatabaseManager() as db:
+            db.set_session(readonly=True)
 
+            with db.cursor() as cursor:
                 # get total number of items so we can estimate the pages
                 cursor.execute('SELECT COUNT(id) FROM items')
                 pages = round(cursor.fetchone()[0] / limit)
@@ -37,8 +36,6 @@ class AllItemsResource:
                 for item in cursor:
                     statistics.append({'id': item['id'], 'views': item['views'], 'downloads': item['downloads']})
 
-        cursor.close()
-
         message = {
             'currentPage': page,
             'totalPages': pages,
@@ -53,6 +50,10 @@ class ItemResource:
     def on_get(self, req, resp, item_id):
         """Handles GET requests"""
 
+        with DatabaseManager() as db:
+            db.set_session(readonly=True)
+
+            with db.cursor() as cursor:
                 cursor = db.cursor()
                 cursor.execute('SELECT views, downloads FROM items WHERE id={}'.format(item_id))
                 if cursor.rowcount == 0:
@@ -71,8 +72,6 @@ class ItemResource:
 
                     resp.media = statistics
 
-        cursor.close()
-
 
 api = application = falcon.API()
 api.add_route('/', RootResource())

dspace_statistics_api/database.py

@@ -7,9 +7,17 @@ import psycopg2
 import psycopg2.extras
 
 
-def database_connection():
-    connection = psycopg2.connect("dbname={} user={} password={} host={} port={}".format(DATABASE_NAME, DATABASE_USER, DATABASE_PASS, DATABASE_HOST, DATABASE_PORT), cursor_factory=psycopg2.extras.DictCursor)
+class DatabaseManager():
+    '''Manage database connection.'''
 
-    return connection
+    def __init__(self):
+        self._connection_uri = 'dbname={} user={} password={} host={} port={}'.format(DATABASE_NAME, DATABASE_USER, DATABASE_PASS, DATABASE_HOST, DATABASE_PORT)
+
+    def __enter__(self):
+        self._connection = psycopg2.connect(self._connection_uri, cursor_factory=psycopg2.extras.DictCursor)
+        return self._connection
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self._connection.close()
 
 # vim: set sw=4 ts=4 expandtab:

dspace_statistics_api/indexer.py

@@ -29,7 +29,7 @@
 # See: https://solrclient.readthedocs.io/en/latest/SolrClient.html
 # See: https://wiki.duraspace.org/display/DSPACE/Solr
 
-from .database import database_connection
+from .database import DatabaseManager
 import json
 import psycopg2.extras
 from .solr import solr_connection
@@ -63,8 +63,8 @@ def index_views():
     results_num_pages = int(results_totalNumFacets / results_per_page)
     results_current_page = 0
 
-    cursor = db.cursor()
-
+    with DatabaseManager() as db:
+        with db.cursor() as cursor:
             # create an empty list to store values for batch insertion
             data = []
 
@@ -97,8 +97,6 @@ def index_views():
 
                 results_current_page += 1
 
-    cursor.close()
-
 
 def index_downloads():
     # get the total number of distinct facets for items with at least 1 download
@@ -123,8 +121,8 @@ def index_downloads():
     results_num_pages = int(results_totalNumFacets / results_per_page)
     results_current_page = 0
 
-    cursor = db.cursor()
-
+    with DatabaseManager() as db:
+        with db.cursor() as cursor:
             # create an empty list to store values for batch insertion
             data = []
 
@@ -157,19 +155,19 @@ def index_downloads():
 
                 results_current_page += 1
 
-    cursor.close()
 
-
-db = database_connection()
 solr = solr_connection()
 
-# create table to store item views and downloads
-cursor = db.cursor()
-cursor.execute('''CREATE TABLE IF NOT EXISTS items
+with DatabaseManager() as db:
+    with db.cursor() as cursor:
+        # create table to store item views and downloads
+        cursor.execute('''CREATE TABLE IF NOT EXISTS items
                   (id INT PRIMARY KEY, views INT DEFAULT 0, downloads INT DEFAULT 0)''')
+
+    # commit the table creation before closing the database connection
+    db.commit()
+
 index_views()
 index_downloads()
 
-db.close()
-
 # vim: set sw=4 ts=4 expandtab: