-
Notifications
You must be signed in to change notification settings - Fork 231
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix spark cluster start mechanism and add extra dev requirements (#986)
- Loading branch information
1 parent
5d90ff9
commit 613fa58
Showing
3 changed files
with
36 additions
and
13 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,19 +1,27 @@ | ||
from multiprocessing import Lock | ||
|
||
import time | ||
import pytest | ||
|
||
_db_start_lock = Lock() | ||
_DB_CLUSTER_STARTED = False | ||
|
||
def _wait_for_databricks_cluster(project): | ||
""" | ||
It takes roughly 3min for the cluster to start, to be safe we'll wait for 5min | ||
""" | ||
for _ in range(60): | ||
try: | ||
project.run_sql("SELECT 1", fetch=True) | ||
return | ||
except Exception: | ||
time.sleep(10) | ||
|
||
raise Exception("Databricks cluster did not start in time") | ||
|
||
|
||
# Running this should prevent tests from needing to be retried because the Databricks cluster isn't available | ||
@pytest.fixture(scope="class", autouse=True) | ||
def start_databricks_cluster(project, request): | ||
global _DB_CLUSTER_STARTED | ||
profile_type = request.config.getoption("--profile") | ||
with _db_start_lock: | ||
if "databricks" in profile_type and not _DB_CLUSTER_STARTED: | ||
print("Starting Databricks cluster") | ||
project.run_sql("SELECT 1") | ||
|
||
_DB_CLUSTER_STARTED = True | ||
if "databricks" in profile_type: | ||
_wait_for_databricks_cluster(project) | ||
|
||
yield 1 |