Make adding/removing databases easier (#4549)

* Make adding/removing databases easier

* Load the database folders without having to edit an index of all the databases

* Remove db descriptions and add some comments

toolset/benchmark/test_types/framework_test_type.py

@@ -1,10 +1,5 @@
 import copy
-import json
 import requests
-import MySQLdb
-import psycopg2
-import pymongo
-import traceback
 
 from colorama import Fore
 from toolset.utils.output_helper import log
@@ -134,81 +129,3 @@ class FrameworkTestType:
         Use before calling parse
         '''
         return copy.copy(self)
-
-    def get_current_world_table(self):
-        '''
-        Return a JSON object containing all 10,000 World items as they currently
-        exist in the database. This is used for verifying that entries in the
-        database have actually changed during an Update verification test.
-        '''
-        database_name = ""
-        results_json = []
-        try:
-            database_name = self.database.lower()
-        except AttributeError:
-            pass
-
-        if database_name == "mysql":
-            try:
-                db = MySQLdb.connect(self.config.database_host,
-                                     "benchmarkdbuser", "benchmarkdbpass",
-                                     "hello_world")
-                cursor = db.cursor()
-                cursor.execute("SELECT * FROM World")
-                results = cursor.fetchall()
-                results_json.append(json.loads(json.dumps(dict(results))))
-                db.close()
-            except Exception:
-                tb = traceback.format_exc()
-                log("ERROR: Unable to load current MySQL World table.",
-                    color=Fore.RED)
-                log(tb)
-        elif database_name == "postgres":
-            try:
-                db = psycopg2.connect(
-                    host=self.config.database_host,
-                    port="5432",
-                    user="benchmarkdbuser",
-                    password="benchmarkdbpass",
-                    database="hello_world")
-                cursor = db.cursor()
-                cursor.execute("SELECT * FROM \"World\"")
-                results = cursor.fetchall()
-                results_json.append(json.loads(json.dumps(dict(results))))
-                cursor = db.cursor()
-                cursor.execute("SELECT * FROM \"world\"")
-                results = cursor.fetchall()
-                results_json.append(json.loads(json.dumps(dict(results))))
-                db.close()
-            except Exception:
-                tb = traceback.format_exc()
-                log("ERROR: Unable to load current Postgres World table.",
-                    color=Fore.RED)
-                log(tb)
-        elif database_name == "mongodb":
-            try:
-                worlds_json = {}
-                print("DATABASE_HOST: %s" % self.config.database_host)
-                connection = pymongo.MongoClient(
-                    host=self.config.database_host)
-                db = connection.hello_world
-                for world in db.world.find():
-                    if "randomNumber" in world:
-                        if "id" in world:
-                            worlds_json[str(int(world["id"]))] = int(
-                                world["randomNumber"])
-                        elif "_id" in world:
-                            worlds_json[str(int(world["_id"]))] = int(
-                                world["randomNumber"])
-                results_json.append(worlds_json)
-                connection.close()
-            except Exception:
-                tb = traceback.format_exc()
-                log("ERROR: Unable to load current MongoDB World table.",
-                    color=Fore.RED)
-                log(tb)
-        else:
-            raise ValueError(
-                "Database: {!s} does not exist".format(database_name))
-
-        return results_json

toolset/benchmark/test_types/verifications.py

@@ -4,6 +4,7 @@ import traceback
 
 from datetime import datetime
 from toolset.utils.output_helper import log
+from toolset.databases import databases
 from time import sleep
 
 def basic_body_verification(body, url, is_json_check=True):
@@ -333,7 +334,7 @@ def verify_query_cases(self, cases, url, check_updates=False):
     # Only load in the World table if we are doing an Update verification
     world_db_before = {}
     if check_updates:
-        world_db_before = self.get_current_world_table()
+        world_db_before = databases[self.database.lower()].get_current_world_table(self.config)
 
     for q, max_infraction in cases:
         case_url = url + q
@@ -358,7 +359,7 @@ def verify_query_cases(self, cases, url, check_updates=False):
             # that only updates 1 item and happens to set its randomNumber to the same value it
             # previously held
             if check_updates and queries >= MAX:
-                world_db_after = self.get_current_world_table()
+                world_db_after = databases[self.database.lower()].get_current_world_table(self.config)
                 problems += verify_updates(world_db_before, world_db_after,
                                            MAX, case_url)
 

toolset/databases/__init__.py

@@ -0,0 +1,24 @@
+import imp
+import re
+
+from colorama import Fore
+from glob import glob
+from toolset.utils.output_helper import log
+
+databases = {}
+db_folders = glob("/FrameworkBenchmarks/toolset/databases/*/")
+
+# Loads all the databases from the folders in this directory
+# and checks to see if they've implemented the required methods
+for folder in db_folders:
+    # regex that grabs the characters between "toolset/database/"
+    # and the final "/" in the db folder string to get the db name
+    db_name = re.findall(r'.+\/(.+)\/$', folder, re.M)[0]
+    db = imp.load_source("Database", "%s%s.py" % (folder, db_name))
+
+    if not hasattr(db.Database, "get_current_world_table")\
+            or not hasattr(db.Database, "test_connection"):
+        log("Database %s does not implement the required methods" + db_name,
+            color=Fore.RED)
+
+    databases[db_name] = db.Database

toolset/databases/mongodb/mongodb.py

@@ -0,0 +1,51 @@
+import pymongo
+import traceback
+
+from colorama import Fore
+from toolset.utils.output_helper import log
+
+class Database:
+
+    @staticmethod
+    def get_current_world_table(config):
+        '''
+        Return a JSON object containing all 10,000 World items as they currently
+        exist in the database. This is used for verifying that entries in the
+        database have actually changed during an Update verification test.
+        '''
+        results_json = []
+
+        try:
+            worlds_json = {}
+            print("DATABASE_HOST: %s" % config.database_host)
+            connection = pymongo.MongoClient(
+                host=config.database_host)
+            db = connection.hello_world
+            for world in db.world.find():
+                if "randomNumber" in world:
+                    if "id" in world:
+                        worlds_json[str(int(world["id"]))] = int(
+                            world["randomNumber"])
+                    elif "_id" in world:
+                        worlds_json[str(int(world["_id"]))] = int(
+                            world["randomNumber"])
+            results_json.append(worlds_json)
+            connection.close()
+        except Exception:
+            tb = traceback.format_exc()
+            log("ERROR: Unable to load current MongoDB World table.",
+                color=Fore.RED)
+            log(tb)
+
+        return results_json
+
+    @staticmethod
+    def test_connection(config):
+        try:
+            connection = pymongo.MongoClient(host=config.database_host)
+            db = connection.hello_world
+            db.world.find()
+            db.close()
+            return True
+        except:
+            return False

toolset/databases/mysql/mysql.py

@@ -0,0 +1,47 @@
+import json
+import MySQLdb
+import traceback
+
+from colorama import Fore
+from toolset.utils.output_helper import log
+
+class Database:
+
+    @staticmethod
+    def get_current_world_table(config):
+        '''
+        Return a JSON object containing all 10,000 World items as they currently
+        exist in the database. This is used for verifying that entries in the
+        database have actually changed during an Update verification test.
+        '''
+        results_json = []
+
+        try:
+            db = MySQLdb.connect(config.database_host,
+                                 "benchmarkdbuser", "benchmarkdbpass",
+                                 "hello_world")
+            cursor = db.cursor()
+            cursor.execute("SELECT * FROM World")
+            results = cursor.fetchall()
+            results_json.append(json.loads(json.dumps(dict(results))))
+            db.close()
+        except Exception:
+            tb = traceback.format_exc()
+            log("ERROR: Unable to load current MySQL World table.",
+                color=Fore.RED)
+            log(tb)
+
+        return results_json
+
+    @staticmethod
+    def test_connection(config):
+        try:
+            db = MySQLdb.connect(config.database_host, "benchmarkdbuser",
+                                 "benchmarkdbpass", "hello_world")
+            cursor = db.cursor()
+            cursor.execute("SELECT 1")
+            cursor.fetchall()
+            db.close()
+            return True
+        except:
+            return False

toolset/databases/postgres/postgres.py

@@ -0,0 +1,58 @@
+import json
+import psycopg2
+import traceback
+
+from colorama import Fore
+from toolset.utils.output_helper import log
+
+class Database:
+
+    @staticmethod
+    def get_current_world_table(config):
+        '''
+        Return a JSON object containing all 10,000 World items as they currently
+        exist in the database. This is used for verifying that entries in the
+        database have actually changed during an Update verification test.
+        '''
+        results_json = []
+
+        try:
+            db = psycopg2.connect(
+                host=config.database_host,
+                port="5432",
+                user="benchmarkdbuser",
+                password="benchmarkdbpass",
+                database="hello_world")
+            cursor = db.cursor()
+            cursor.execute("SELECT * FROM \"World\"")
+            results = cursor.fetchall()
+            results_json.append(json.loads(json.dumps(dict(results))))
+            cursor = db.cursor()
+            cursor.execute("SELECT * FROM \"world\"")
+            results = cursor.fetchall()
+            results_json.append(json.loads(json.dumps(dict(results))))
+            db.close()
+        except Exception:
+            tb = traceback.format_exc()
+            log("ERROR: Unable to load current Postgres World table.",
+                color=Fore.RED)
+            log(tb)
+
+        return results_json
+
+    @staticmethod
+    def test_connection(config):
+        try:
+            db = psycopg2.connect(
+                host=config.database_host,
+                port="5432",
+                user="benchmarkdbuser",
+                password="benchmarkdbpass",
+                database="hello_world")
+            cursor = db.cursor()
+            cursor.execute("SELECT 1")
+            cursor.fetchall()
+            db.close()
+            return True
+        except:
+            return False

toolset/utils/database_helper.py

@@ -1,40 +0,0 @@
-import MySQLdb
-import psycopg2
-import pymongo
-
-
-def test_database(config, database_name):
-    if database_name == "mysql":
-        try:
-            db = MySQLdb.connect(config.database_host, "benchmarkdbuser",
-                                 "benchmarkdbpass", "hello_world")
-            cursor = db.cursor()
-            cursor.execute("SELECT 1")
-            cursor.fetchall()
-            db.close()
-        except:
-            return False
-    elif database_name == "postgres":
-        try:
-            db = psycopg2.connect(
-                host=config.database_host,
-                port="5432",
-                user="benchmarkdbuser",
-                password="benchmarkdbpass",
-                database="hello_world")
-            cursor = db.cursor()
-            cursor.execute("SELECT 1")
-            cursor.fetchall()
-            db.close()
-        except:
-            return False
-    elif database_name == "mongodb":
-        try:
-            connection = pymongo.MongoClient(host=config.database_host)
-            db = connection.hello_world
-            db.world.find()
-            db.close()
-        except:
-            return False
-
-    return True

toolset/utils/docker_helper.py

@@ -9,7 +9,7 @@ from threading import Thread
 from colorama import Fore, Style
 
 from toolset.utils.output_helper import log
-from toolset.utils.database_helper import test_database
+from toolset.databases import databases
 
 from psutil import virtual_memory
 
@@ -346,7 +346,7 @@ class DockerHelper:
         while not database_ready and slept < max_sleep:
             time.sleep(1)
             slept += 1
-            database_ready = test_database(self.benchmarker.config, database)
+            database_ready = databases[database].test_connection(self.benchmarker.config)
 
         if not database_ready:
             log("Database was not ready after startup", prefix=log_prefix)

toolset/utils/metadata.py

@@ -3,20 +3,16 @@ import glob
 import json
 
 from collections import OrderedDict
-
-from toolset.utils.output_helper import log
 from colorama import Fore
+from toolset.databases import databases
+from toolset.utils.output_helper import log
 
 
 class Metadata:
 
-    supported_dbs = [
-        ('MySQL',
-         'One of the most popular databases around the web and in TFB'),
-        ('Postgres',
-         'An advanced SQL database with a larger feature set than MySQL'),
-        ('MongoDB', 'A popular document-store database')
-    ]
+    supported_dbs = []
+    for name in databases:
+        supported_dbs.append((name, '...'))
 
     def __init__(self, benchmarker=None):
         self.benchmarker = benchmarker

toolset/utils/scaffolding.py

@@ -205,7 +205,7 @@ class Scaffolding:
         prompt = "Database ["
         options = []
         for db in Metadata.supported_dbs:
-            print("  {!s}) {!s}: {!s}".format(i, db[0], db[1]))
+            print("  {!s}) {!s}".format(i, db[0]))
             prompt += "{!s}/".format(i)
             options.append(db[0])
             i += 1