def test_exec_failure(self, db_mock_class):
"""
Test the execute function in case where the run failed.
"""
run = {
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
}
op = DatabricksSubmitRunOperator(task_id=TASK_ID, json=run)
db_mock = db_mock_class.return_value
db_mock.submit_run.return_value = 1
db_mock.get_run_state.return_value = RunState('TERMINATED', 'FAILED',
'')
with self.assertRaises(AirflowException):
op.execute(None)
expected = databricks_operator._deep_string_coerce({
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
'run_name': TASK_ID,
})
db_mock_class.assert_called_once_with(
DEFAULT_CONN_ID,
retry_limit=op.databricks_retry_limit,
retry_delay=op.databricks_retry_delay)
db_mock.submit_run.assert_called_once_with(expected)
db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
db_mock.get_run_state.assert_called_once_with(RUN_ID)
self.assertEqual(RUN_ID, op.run_id)
def test_exec_success(self, db_mock_class):
"""
Test the execute function in case where the run is successful.
"""
run = {
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
}
op = DatabricksSubmitRunOperator(task_id=TASK_ID, json=run)
db_mock = db_mock_class.return_value
db_mock.submit_run.return_value = 1
db_mock.get_run_state.return_value = RunState('TERMINATED', 'SUCCESS',
'')
op.execute(None)
expected = databricks_operator._deep_string_coerce({
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
'run_name': TASK_ID
})
db_mock_class.assert_called_once_with(
DEFAULT_CONN_ID,
retry_limit=op.databricks_retry_limit,
retry_delay=op.databricks_retry_delay)
db_mock.submit_run.assert_called_once_with(expected)
db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
db_mock.get_run_state.assert_called_once_with(RUN_ID)
assert RUN_ID == op.run_id
def test_init_with_bad_type(self):
json = {'test': datetime.now()}
# Looks a bit weird since we have to escape regex reserved symbols.
exception_message = r'Type \<(type|class) \'datetime.datetime\'\> used ' + \
r'for parameter json\[test\] is not a number or a string'
with self.assertRaisesRegex(AirflowException, exception_message):
DatabricksSubmitRunOperator(task_id=TASK_ID, json=json)
def test_init_with_specified_run_name(self):
"""
Test the initializer with a specified run_name.
"""
json = {'new_cluster': NEW_CLUSTER, 'notebook_task': NOTEBOOK_TASK, 'run_name': RUN_NAME}
op = DatabricksSubmitRunOperator(task_id=TASK_ID, json=json)
expected = databricks_operator._deep_string_coerce(
{'new_cluster': NEW_CLUSTER, 'notebook_task': NOTEBOOK_TASK, 'run_name': RUN_NAME}
)
self.assertDictEqual(expected, op.json)
def test_on_kill(self, db_mock_class):
run = {
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
}
op = DatabricksSubmitRunOperator(task_id=TASK_ID, json=run)
db_mock = db_mock_class.return_value
op.run_id = RUN_ID
op.on_kill()
db_mock.cancel_run.assert_called_once_with(RUN_ID)
def test_init_with_spark_python_task_named_parameters(self):
"""
Test the initializer with the named parameters.
"""
op = DatabricksSubmitRunOperator(
task_id=TASK_ID, new_cluster=NEW_CLUSTER, spark_python_task=SPARK_PYTHON_TASK
)
expected = databricks_operator._deep_string_coerce(
{'new_cluster': NEW_CLUSTER, 'spark_python_task': SPARK_PYTHON_TASK, 'run_name': TASK_ID}
)
self.assertDictEqual(expected, op.json)
def test_init_with_json(self):
"""
Test the initializer with json data.
"""
json = {'new_cluster': NEW_CLUSTER, 'notebook_task': NOTEBOOK_TASK}
op = DatabricksSubmitRunOperator(task_id=TASK_ID, json=json)
expected = databricks_operator._deep_string_coerce({
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
'run_name': TASK_ID
})
assert expected == op.json
def test_init_with_named_parameters(self):
"""
Test the initializer with the named parameters.
"""
op = DatabricksSubmitRunOperator(task_id=TASK_ID,
new_cluster=NEW_CLUSTER,
notebook_task=NOTEBOOK_TASK)
expected = databricks_operator._deep_string_coerce({
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
'run_name': TASK_ID
})
self.assertDictEqual(expected, op.json)
def test_init_with_templating(self):
json = {
'new_cluster': NEW_CLUSTER,
'notebook_task': TEMPLATED_NOTEBOOK_TASK,
}
dag = DAG('test', start_date=datetime.now())
op = DatabricksSubmitRunOperator(dag=dag, task_id=TASK_ID, json=json)
op.render_template_fields(context={'ds': DATE})
expected = databricks_operator._deep_string_coerce({
'new_cluster': NEW_CLUSTER,
'notebook_task': RENDERED_TEMPLATED_NOTEBOOK_TASK,
'run_name': TASK_ID,
})
self.assertDictEqual(expected, op.json)
def test_init_with_merging(self):
"""
Test the initializer when json and other named parameters are both
provided. The named parameters should override top level keys in the
json dict.
"""
override_new_cluster = {'workers': 999}
json = {
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
}
op = DatabricksSubmitRunOperator(task_id=TASK_ID,
json=json,
new_cluster=override_new_cluster)
expected = databricks_operator._deep_string_coerce({
'new_cluster': override_new_cluster,
'notebook_task': NOTEBOOK_TASK,
'run_name': TASK_ID,
})
self.assertDictEqual(expected, op.json)
new_cluster = {
'spark_version': '2.1.0-db3-scala2.11',
'node_type_id': 'r3.xlarge',
'aws_attributes': {
'availability': 'ON_DEMAND'
},
'num_workers': 8
}
notebook_task_params = {
'new_cluster': new_cluster,
'notebook_task': {
'notebook_path': '/Users/[email protected]/PrepareData',
},
}
# Example of using the JSON parameter to initialize the operator.
notebook_task = DatabricksSubmitRunOperator(task_id='notebook_task',
json=notebook_task_params)
# Example of using the named parameters of DatabricksSubmitRunOperator
# to initialize the operator.
spark_jar_task = DatabricksSubmitRunOperator(
task_id='spark_jar_task',
new_cluster=new_cluster,
spark_jar_task={'main_class_name': 'com.example.ProcessData'},
libraries=[{
'jar': 'dbfs:/lib/etl-0.1.jar'
}])
notebook_task >> spark_jar_task
) as dag:
new_cluster = {
'spark_version': '6.4.x-scala2.11',
'node_type_id': 'r3.xlarge',
'aws_attributes': {'availability': 'ON_DEMAND'},
'num_workers': 2,
}
notebook_task_params1 = {
'new_cluster': new_cluster,
'notebook_task': {
'notebook_path': '/Users/[email protected]/temp/nb1',
},
}
# Example of using the JSON parameter to initialize the operator.
notebook_task1 = DatabricksSubmitRunOperator(task_id='notebook_task1', json=notebook_task_params1)
notebook_task_params2 = {
'new_cluster': new_cluster,
'notebook_task': {
'notebook_path': '/Users/[email protected]/temp/nb2',
},
}
# Example of using the JSON parameter to initialize the operator.
notebook_task2 = DatabricksSubmitRunOperator(task_id='notebook_task2', json=notebook_task_params2)
notebook_task_params3 = {
'new_cluster': new_cluster,
'notebook_task': {
'notebook_path': '/Users/[email protected]/temp/nb3',
},
with DAG(
'orchestration_good_practices',
start_date=datetime(2021, 1, 1),
schedule_interval='@daily',
catchup=False,
default_args={
'owner': 'airflow',
'email_on_failure': False,
'retries': 1,
'retry_delay': timedelta(minutes=1)
},
) as dag:
opr_refresh_mat_view = PostgresOperator(
task_id='refresh_mat_view',
postgres_conn_id='postgres_default',
sql='REFRESH MATERIALIZED VIEW example_view;',
)
opr_submit_run = DatabricksSubmitRunOperator(
task_id='submit_run',
databricks_conn_id='databricks',
new_cluster=new_cluster,
notebook_task=notebook_task)
opr_run_now = DatabricksRunNowOperator(task_id='run_now',
databricks_conn_id='databricks',
job_id=5,
notebook_params=notebook_params)
opr_refresh_mat_view >> opr_submit_run >> opr_run_now
update_repo = DatabricksReposUpdateOperator(task_id='update_repo',
repo_path=repo_path,
branch="releases")
# [END howto_operator_databricks_repo_update]
notebook_task_params = {
'new_cluster': {
'spark_version': '9.1.x-scala2.12',
'node_type_id': 'r3.xlarge',
'aws_attributes': {
'availability': 'ON_DEMAND'
},
'num_workers': 8,
},
'notebook_task': {
'notebook_path': f'{repo_path}/PrepareData',
},
}
notebook_task = DatabricksSubmitRunOperator(task_id='notebook_task',
json=notebook_task_params)
# [START howto_operator_databricks_repo_delete]
# Example of deleting a Databricks Repo
repo_path = "/Repos/[email protected]/demo-repo"
delete_repo = DatabricksReposDeleteOperator(task_id='delete_repo',
repo_path=repo_path)
# [END howto_operator_databricks_repo_delete]
(create_repo >> update_repo >> notebook_task >> delete_repo)
"depends_on_past": False,
"email_on_failure": False,
"email_on_retry": False,
"retries": 1,
"retry_delay": timedelta(minutes=2),
}
with DAG(
"databricks_dag",
start_date=datetime(2021, 1, 1),
schedule_interval="@daily",
catchup=False,
default_args=default_args,
) as dag:
opr_submit_run = DatabricksSubmitRunOperator(
task_id="submit_run",
databricks_conn_id="databricks",
new_cluster=new_cluster,
notebook_task=notebook_task,
)
opr_run_now = DatabricksRunNowOperator(
task_id="run_now",
databricks_conn_id="databricks",
job_id=5,
notebook_params=notebook_params,
)
opr_submit_run >> opr_run_now