Skip to content

Commit

Permalink
Added incompatible Databricks Runtime check for Airflow operators
Browse files Browse the repository at this point in the history
Fixes #3
  • Loading branch information
nfx committed Mar 15, 2024
1 parent 8ab5ad8 commit 966746f
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 0 deletions.
20 changes: 20 additions & 0 deletions src/databricks/labs/pylint/airflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,11 @@ class AirflowChecker(BaseChecker):
"missing-data-security-mode",
"new_cluster is missing data_security_mode",
),
"E9698": (
"%s cluster has unsupported runtime: %s",
"unsupported-runtime",
"new_cluster has unsupported runtime",
),
}

def visit_call(self, node: astroid.Call):
Expand All @@ -30,6 +35,21 @@ def visit_call(self, node: astroid.Call):
def _check_new_cluster(self, key: str, new_cluster: dict[str, Any], node: astroid.NodeNG):
if "data_security_mode" not in new_cluster:
self.add_message("missing-data-security-mode", node=node, args=(key,))
if "spark_version" in new_cluster and not self._is_supported(new_cluster["spark_version"]):
self.add_message("unsupported-runtime", node=node, args=(key, new_cluster["spark_version"]))

@staticmethod
def _is_supported(spark_version: str):
try:
split = spark_version.split("-")
if len(split) < 2:
return False
digits = split[0].split(".")
if len(digits) < 2:
return False
return (int(digits[0]), int(digits[1])) >= (11, 3)
except ValueError:
return False

def _check_tasks(self, tasks: list[dict[str, Any]], node: astroid.NodeNG):
for task in tasks:
Expand Down
1 change: 1 addition & 0 deletions tests/test_airflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ def test_missing_data_security_mode_in_submit_run_clusters(lint_with):
task_id="notebook_run", new_cluster=new_cluster, notebook_task=notebook_task
)"""
)
assert "[unsupported-runtime] ephemeral cluster has unsupported runtime: 10.1.x-scala2.12" in messages
assert (
"[missing-data-security-mode] ephemeral cluster missing 'data_security_mode' "
"required for Unity Catalog compatibility"
Expand Down

0 comments on commit 966746f

Please sign in to comment.