Merge pull request #226 from opensafely/fix-presto-tests

Fix flaky Presto tests

cohortextractor/presto_utils.py

@@ -5,10 +5,6 @@
 import requests
 
 
-class PrestoNoActiveNodesError(Exception):
-    pass
-
-
 def presto_connection_from_url(url):
     return ConnectionProxy(
         prestodb.dbapi.connect(**presto_connection_params_from_url(url))
@@ -40,24 +36,24 @@ def presto_connection_params_from_url(url):
     return connection_params
 
 
-def wait_for_presto_to_be_ready(url, timeout):
+def wait_for_presto_to_be_ready(url, test_query, timeout):
+    """
+    Waits for Presto to be ready to execute queries by repeatedly attempting to
+    connect and run `test_query`, raising the last received error after
+    `timeout` seconds
+    """
     connection_params = presto_connection_params_from_url(url)
     start = time.time()
     while True:
         try:
             connection = prestodb.dbapi.connect(**connection_params)
             cursor = connection.cursor()
-            cursor.execute(
-                "SELECT COUNT(*) FROM system.runtime.nodes WHERE state = 'active'"
-            )
-            count = cursor.fetchone()[0]
-            if count == 0:
-                raise PrestoNoActiveNodesError("No active nodes found")
+            cursor.execute(test_query)
+            cursor.fetchall()
             break
         except (
             prestodb.exceptions.PrestoQueryError,
             requests.exceptions.ConnectionError,
-            PrestoNoActiveNodesError,
         ):
             if time.time() - start < timeout:
                 time.sleep(1)

run.sh

@@ -35,7 +35,7 @@ exec docker-compose run \
   --volume "$PWD:/workspace" \
   -e PYENV_VERSION=3.7.8 \
   -e TPP_DATABASE_URL='mssql://SA:Your_password123!@mssql:1433/Test_OpenCorona' \
-  -e EMIS_DATABASE_URL=presto://presto:8080/mssql/dbo \
-  -e EMIS_DATASOURCE_DATABASE_URL='mssql://SA:Your_password123!@mssql:1433/Test_EMIS' \
+  -e ACME_DATABASE_URL=presto://presto:8080/mssql/dbo \
+  -e ACME_DATASOURCE_DATABASE_URL='mssql://SA:Your_password123!@mssql:1433/Test_ACME' \
   app \
     --login -- "$@"

tests/acme_backend_setup.py

@@ -60,7 +60,17 @@ def make_engine():
                 time.sleep(1)
             else:
                 raise
-    wait_for_presto_to_be_ready(os.environ["ACME_DATABASE_URL"], timeout)
+    wait_for_presto_to_be_ready(
+        os.environ["ACME_DATABASE_URL"],
+        # Presto will show active nodes in its `system.runtime.nodes` table but
+        # then throw a "no nodes available" error if you try to execute a query
+        # which needs to touch the MSSQL instance. So to properly confirm that
+        # Presto is ready we need a query which forces it to connect to MSSQL,
+        # but ideally one which doesn't depend on any particular configuration
+        # having been done first. The below seems to do the trick.
+        "SELECT 1 FROM sys.tables",
+        timeout,
+    )
     return engine