From a00c2f42e608d52848334a882209f1c01236d91d Mon Sep 17 00:00:00 2001
From: David Moreau Simard <dmsimard@redhat.com>
Date: Tue, 19 Nov 2019 10:49:51 -0500
Subject: [PATCH] Improve database engine error handling when starting the
 server

This will check if the database driver library is installed before
attempting to run the server.

Fixes: https://github.com/ansible-community/ara/issues/63
Change-Id: I42a7eca4909a560aba95d725021f4db1f9682428
---
 ara/server/__main__.py  | 14 +++++++++++++-
 ara/setup/exceptions.py | 14 +++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/ara/server/__main__.py b/ara/server/__main__.py
index 2fd6ac12..e955fbc5 100644
--- a/ara/server/__main__.py
+++ b/ara/server/__main__.py
@@ -21,7 +21,7 @@ import sys
 
 from django.conf import settings
 
-from ara.setup.exceptions import MissingDjangoException
+from ara.setup.exceptions import MissingDjangoException, MissingMysqlclientException, MissingPsycopgException
 
 
 def main():
@@ -32,6 +32,18 @@ def main():
     except ImportError as e:
         raise MissingDjangoException from e
 
+    if settings.DATABASE_ENGINE == "django.db.backends.postgresql":
+        try:
+            import psycopg2  # noqa
+        except ImportError as e:
+            raise MissingPsycopgException from e
+
+    if settings.DATABASE_ENGINE == "django.db.backends.mysql":
+        try:
+            import MySQLdb  # noqa
+        except ImportError as e:
+            raise MissingMysqlclientException from e
+
     execute_from_command_line(sys.argv)
     print("[ara] Using settings file: %s" % settings.ARA_SETTINGS)
 
diff --git a/ara/setup/exceptions.py b/ara/setup/exceptions.py
index b8815ba6..4ea56645 100644
--- a/ara/setup/exceptions.py
+++ b/ara/setup/exceptions.py
@@ -18,5 +18,17 @@
 
 class MissingDjangoException(Exception):
     def __init__(self):
-        exc = "Unable to import Django: the server dependencies can be installed with 'pip install ara[server]'"
+        exc = "The server dependencies must be installed to record data offline or run the API server."
+        super().__init__(exc)
+
+
+class MissingPsycopgException(Exception):
+    def __init__(self):
+        exc = "The psycopg2 python library must be installed in order to use the PostgreSQL database engine."
+        super().__init__(exc)
+
+
+class MissingMysqlclientException(Exception):
+    def __init__(self):
+        exc = "The mysqlclient python library must be installed in order to use the MySQL database engine."
         super().__init__(exc)