improved error messages, added color-coding

Author: dmalan

src/cs50/sql.py

@@ -2,10 +2,12 @@
 import decimal
 import importlib
 import logging
+import os
 import re
 import sqlalchemy
 import sqlparse
 import sys
+import termcolor
 import warnings
 
 
@@ -22,12 +24,52 @@ def __init__(self, url, **kwargs):
         http://docs.sqlalchemy.org/en/latest/dialects/index.html
         """
 
+        # Require that file already exist for SQLite
+        matches = re.search(r"^sqlite:///(.+)$", url)
+        if matches:
+            if not os.path.exists(matches.group(1)):
+                raise RuntimeError("does not exist: {}".format(matches.group(1)))
+            if not os.path.isfile(matches.group(1)):
+                raise RuntimeError("not a file: {}".format(matches.group(1)))
+
+        # Create engine, raising exception if back end's module not installed
+        self.engine = sqlalchemy.create_engine(url, **kwargs)
+
         # Log statements to standard error
         logging.basicConfig(level=logging.DEBUG)
         self.logger = logging.getLogger("cs50")
 
-        # Create engine, raising exception if back end's module not installed
-        self.engine = sqlalchemy.create_engine(url, **kwargs)
+        # Test database
+        try:
+            self.logger.disabled = True
+            self.execute("SELECT 1")
+        except sqlalchemy.exc.OperationalError as e:
+            e = RuntimeError(self._parse(e))
+            e.__cause__ = None
+            raise e
+        else:
+            self.logger.disabled = False
+
+    def _parse(self, e):
+        """Parses an exception, returns its message."""
+
+        # MySQL
+        matches = re.search(r"^\(_mysql_exceptions\.OperationalError\) \(\d+, \"(.+)\"\)$", str(e))
+        if matches:
+            return matches.group(1)
+
+        # PostgreSQL
+        matches = re.search(r"^\((psycopg2\.OperationalError)\) (.+)$", str(e))
+        if matches:
+            return matches.group(1)
+
+        # SQLite
+        matches = re.search(r"^\(sqlite3\.OperationalError\) (.+)$", str(e))
+        if matches:
+            return matches.group(1)
+
+        # Default
+        return str(e)
 
     def execute(self, text, **params):
         """
@@ -119,12 +161,12 @@ def process(value):
             # http://docs.sqlalchemy.org/en/latest/faq/sqlexpressions.html#how-do-i-render-sql-expressions-as-strings-possibly-with-bound-parameters-inlined
             statement = str(statement.compile(compile_kwargs={"literal_binds": True}))
 
+            # Statement for logging
+            log = re.sub(r"\n\s*", " ", sqlparse.format(statement, reindent=True))
+
             # Execute statement
             result = self.engine.execute(statement)
 
-            # Log statement
-            self.logger.debug(re.sub(r"\n\s*", " ", sqlparse.format(statement, reindent=True)))
-
             # If SELECT (or INSERT with RETURNING), return result set as list of dict objects
             if re.search(r"^\s*SELECT", statement, re.I):
 
@@ -135,23 +177,36 @@ def process(value):
                     for column in row:
                         if isinstance(row[column], decimal.Decimal):
                             row[column] = float(row[column])
-                return rows
+                ret = rows
 
             # If INSERT, return primary key value for a newly inserted row
             elif re.search(r"^\s*INSERT", statement, re.I):
                 if self.engine.url.get_backend_name() in ["postgres", "postgresql"]:
                     result = self.engine.execute(sqlalchemy.text("SELECT LASTVAL()"))
-                    return result.first()[0]
+                    ret = result.first()[0]
                 else:
-                    return result.lastrowid
+                    ret = result.lastrowid
 
             # If DELETE or UPDATE, return number of rows matched
             elif re.search(r"^\s*(?:DELETE|UPDATE)", statement, re.I):
-                return result.rowcount
+                ret = result.rowcount
 
             # If some other statement, return True unless exception
-            return True
+            ret = True
 
         # If constraint violated, return None
         except sqlalchemy.exc.IntegrityError:
+            self.logger.debug(termcolor.colored(log, "yellow"))
             return None
+
+        # If user errror
+        except sqlalchemy.exc.OperationalError as e:
+            self.logger.debug(termcolor.colored(log, "red"))
+            e = RuntimeError(self._parse(e))
+            e.__cause__ = None
+            raise e
+
+        # Return value
+        else:
+            self.logger.debug(termcolor.colored(log, "green"))
+            return ret