Making sqla checkout check cross-db

apache · Sep 17, 2016 · a64ee9b · a64ee9b
1 parent ca351cf
commit a64ee9b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/caravel/__init__.py b/caravel/__init__.py
@@ -43,7 +43,11 @@ def checkout(dbapi_con, con_record, con_proxy):
     """
     try:
         try:
-            dbapi_con.ping(False)
+            if hasattr(dbapi_con, 'ping'):
+                dbapi_con.ping(False)
+            else:
+                cursor = dbapi_con.cursor()
+                cursor.execute("SELECT 1")
         except TypeError:
             app.logger.debug('MySQL connection died. Restoring...')
             dbapi_con.ping()

diff --git a/tests/celery_tests.py b/tests/celery_tests.py
@@ -226,7 +226,8 @@ def test_run_async_query(self):
         sql_where = "SELECT name FROM ab_role WHERE name='Admin'"
         result1 = self.run_sql(
             1, sql_where, async='true', tmp_table='tmp_async_1', cta='true')
-        self.assertEqual(QueryStatus.PENDING, result1['query']['state'])
+        assert result1['query']['state'] in (
+            QueryStatus.PENDING, QueryStatus.RUNNING, QueryStatus.SUCCESS)
 
         time.sleep(1)