Fix #571 -- Improve database engine detection.

3 files changed, 16 insertions, 23 deletions

diff --git a/debug_toolbar/panels/sql/panel.py b/debug_toolbar/panels/sql/panel.py
index f6ce954..5531964 100644
--- a/debug_toolbar/panels/sql/panel.py
+++ b/debug_toolbar/panels/sql/panel.py
@@ -15,8 +15,8 @@ from debug_toolbar.panels.sql.utils import reformat_sql, contrasting_color_gener
 from debug_toolbar.panels.sql.tracking import wrap_cursor, unwrap_cursor
 
 
-def get_isolation_level_display(engine, level):
-    if engine == 'psycopg2':
+def get_isolation_level_display(vendor, level):
+    if vendor == 'postgresql':
         import psycopg2.extensions
         choices = {
             psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT: _("Autocommit"),
@@ -26,12 +26,12 @@ def get_isolation_level_display(engine, level):
             psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE: _("Serializable"),
         }
     else:
-        raise ValueError(engine)
+        raise ValueError(vendor)
     return choices.get(level)
 
 
-def get_transaction_status_display(engine, level):
-    if engine == 'psycopg2':
+def get_transaction_status_display(vendor, level):
+    if vendor == 'postgresql':
         import psycopg2.extensions
         choices = {
             psycopg2.extensions.TRANSACTION_STATUS_IDLE: _("Idle"),
@@ -41,7 +41,7 @@ def get_transaction_status_display(engine, level):
             psycopg2.extensions.TRANSACTION_STATUS_UNKNOWN: _("Unknown"),
         }
     else:
-        raise ValueError(engine)
+        raise ValueError(vendor)
     return choices.get(level)
 
 
@@ -67,11 +67,10 @@ class SQLPanel(Panel):
         if not conn:
             return
 
-        engine = conn.__class__.__module__.split('.', 1)[0]
-        if engine == 'psycopg2':
+        if conn.vendor == 'postgresql':
             cur_status = conn.get_transaction_status()
         else:
-            raise ValueError(engine)
+            raise ValueError(conn.vendor)
 
         last_status = self._transaction_status.get(alias)
         self._transaction_status[alias] = cur_status
@@ -175,10 +174,10 @@ class SQLPanel(Panel):
 
                 query['alias'] = alias
                 if 'iso_level' in query:
-                    query['iso_level'] = get_isolation_level_display(query['engine'],
+                    query['iso_level'] = get_isolation_level_display(query['vendor'],
                                                                      query['iso_level'])
                 if 'trans_status' in query:
-                    query['trans_status'] = get_transaction_status_display(query['engine'],
+                    query['trans_status'] = get_transaction_status_display(query['vendor'],
                                                                            query['trans_status'])
 
                 query['form'] = SQLSelectForm(auto_id=None, initial=copy(query))
diff --git a/debug_toolbar/panels/sql/tracking.py b/debug_toolbar/panels/sql/tracking.py
index 68c77e2..b6a787d 100644
--- a/debug_toolbar/panels/sql/tracking.py
+++ b/debug_toolbar/panels/sql/tracking.py
@@ -131,14 +131,10 @@ class NormalCursorWrapper(object):
 
             alias = getattr(self.db, 'alias', 'default')
             conn = self.db.connection
-            # HACK: avoid imports
-            if conn:
-                engine = conn.__class__.__module__.split('.', 1)[0]
-            else:
-                engine = 'unknown'
+            vendor = getattr(conn, 'vendor', 'unknown')
 
             params = {
-                'engine': engine,
+                'vendor': vendor,
                 'alias': alias,
                 'sql': self.db.ops.last_executed_query(
                     self.cursor, sql, self._quote_params(params)),
@@ -153,7 +149,7 @@ class NormalCursorWrapper(object):
                 'template_info': template_info,
             }
 
-            if engine == 'psycopg2':
+            if vendor == 'postgresql':
                 # If an erroneous query was ran on the connection, it might
                 # be in a state where checking isolation_level raises an
                 # exception.
diff --git a/debug_toolbar/panels/sql/views.py b/debug_toolbar/panels/sql/views.py
index 05ad74f..d8c94a0 100644
--- a/debug_toolbar/panels/sql/views.py
+++ b/debug_toolbar/panels/sql/views.py
@@ -39,17 +39,15 @@ def sql_explain(request):
     if form.is_valid():
         sql = form.cleaned_data['raw_sql']
         params = form.cleaned_data['params']
+        vendor = form.connection.vendor
         cursor = form.cursor
 
-        conn = form.connection
-        engine = conn.__class__.__module__.split('.', 1)[0]
-
-        if engine == "sqlite3":
+        if vendor == 'sqlite':
             # SQLite's EXPLAIN dumps the low-level opcodes generated for a query;
             # EXPLAIN QUERY PLAN dumps a more human-readable summary
             # See http://www.sqlite.org/lang_explain.html for details
             cursor.execute("EXPLAIN QUERY PLAN %s" % (sql,), params)
-        elif engine == "psycopg2":
+        elif vendor == 'postgresql':
             cursor.execute("EXPLAIN ANALYZE %s" % (sql,), params)
         else:
             cursor.execute("EXPLAIN %s" % (sql,), params)