Add support for dbt tasks (#520)

Co-authored-by: Pieter Noordhuis <[email protected]>

Author: kzidane

databricks_cli/sdk/service.py

@@ -57,6 +57,7 @@ def create_job(
         job_clusters=None,
         tags=None,
         format=None,
+        dbt_task=None,
     ):
         _data = {}
         if name is not None:
@@ -125,6 +126,10 @@ def create_job(
             _data['tags'] = tags
         if format is not None:
             _data['format'] = format
+        if dbt_task is not None:
+            _data['dbt_task'] = dbt_task
+            if not isinstance(dbt_task, dict):
+                raise TypeError('Expected databricks.DbtTask() or dict for field dbt_task')
         return self.client.perform_query(
             'POST', '/jobs/create', data=_data, headers=headers, version=version
         )
@@ -146,6 +151,7 @@ def submit_run(
         idempotency_token=None,
         job_clusters=None,
         git_source=None,
+        dbt_task=None,
     ):
         _data = {}
         if run_name is not None:
@@ -194,6 +200,10 @@ def submit_run(
             _data['git_source'] = git_source
             if not isinstance(git_source, dict):
                 raise TypeError('Expected databricks.GitSource() or dict for field git_source')
+        if dbt_task is not None:
+            _data['dbt_task'] = dbt_task
+            if not isinstance(dbt_task, dict):
+                raise TypeError('Expected databricks.DbtTask() or dict for field dbt_task')
         return self.client.perform_query(
             'POST', '/jobs/runs/submit', data=_data, headers=headers, version=version
         )
@@ -253,6 +263,7 @@ def run_now(
         idempotency_token=None,
         headers=None,
         version=None,
+        dbt_commands=None,
     ):
         _data = {}
         if job_id is not None:
@@ -269,6 +280,8 @@ def run_now(
             _data['python_named_params'] = python_named_params
         if idempotency_token is not None:
             _data['idempotency_token'] = idempotency_token
+        if dbt_commands is not None:
+            _data['dbt_commands'] = dbt_commands
         return self.client.perform_query(
             'POST', '/jobs/run-now', data=_data, headers=headers, version=version
         )
@@ -285,6 +298,7 @@ def repair(
         python_named_params=None,
         headers=None,
         version=None,
+        dbt_commands=None,
     ):
         _data = {}
         if run_id is not None:
@@ -303,6 +317,8 @@ def repair(
             _data['spark_submit_params'] = spark_submit_params
         if python_named_params is not None:
             _data['python_named_params'] = python_named_params
+        if dbt_commands is not None:
+            _data['dbt_commands'] = dbt_commands
         return self.client.perform_query(
             'POST', '/jobs/runs/repair', data=_data, headers=headers, version=version
         )

tests/sdk/test_service.py

@@ -130,6 +130,36 @@ def test_create_job(jobs_service):
     jobs_service.client.perform_query.assert_called_with('POST', '/jobs/create', data={'tasks': tasks}, headers=None, version='2.1')
 
 
+@provide_conf
+def test_create_dbt_task(jobs_service):
+    git_source  = {
+        'git_provider': 'github',
+        'git_url': 'https://github.com/foo/bar',
+        'git_branch': 'main'
+    }
+
+    tasks = [
+      {
+        'task_key': 'dbt',
+        'dbt_task': {
+            'commands': ['dbt test']
+        }
+      }
+    ]
+
+    jobs_service.create_job(git_source=git_source, tasks=tasks)
+    jobs_service.client.perform_query.assert_called_with('POST', '/jobs/create', data={'git_source': git_source, 'tasks': tasks}, headers=None, version=None)
+
+
+@provide_conf
+def test_run_now_dbt_task(jobs_service):
+    job_id = 1337
+    dbt_commands = ['dbt test', 'dbt deps']
+
+    jobs_service.run_now(job_id=job_id, dbt_commands=dbt_commands)
+    jobs_service.client.perform_query.assert_called_with('POST', '/jobs/run-now', data={'job_id': job_id, 'dbt_commands': dbt_commands}, headers=None, version=None)
+
+
 @provide_conf
 def test_create_job_invalid_types(jobs_service):
     with pytest.raises(TypeError, match='new_cluster'):
@@ -140,7 +170,10 @@ def test_create_job_invalid_types(jobs_service):
 
     with pytest.raises(TypeError, match='schedule'):
         jobs_service.create_job(schedule=[])
-        
+
+    with pytest.raises(TypeError, match='git_source'):
+        jobs_service.create_job(git_source=[])
+
     with pytest.raises(TypeError, match='notebook_task'):
         jobs_service.create_job(notebook_task=[])
 
@@ -153,6 +186,9 @@ def test_create_job_invalid_types(jobs_service):
     with pytest.raises(TypeError, match='spark_submit_task'):
         jobs_service.create_job(spark_submit_task=[])
 
+    with pytest.raises(TypeError, match='dbt_task'):
+        jobs_service.create_job(dbt_task=[])
+
 
 @provide_conf
 def test_submit_run_invalid_types(jobs_service):
@@ -164,7 +200,10 @@ def test_submit_run_invalid_types(jobs_service):
 
     with pytest.raises(TypeError, match='schedule'):
         jobs_service.submit_run(schedule=[])
-        
+
+    with pytest.raises(TypeError, match='git_source'):
+        jobs_service.submit_run(git_source=[])
+
     with pytest.raises(TypeError, match='notebook_task'):
         jobs_service.submit_run(notebook_task=[])
 
@@ -176,3 +215,6 @@ def test_submit_run_invalid_types(jobs_service):
 
     with pytest.raises(TypeError, match='spark_submit_task'):
         jobs_service.submit_run(spark_submit_task=[])
+
+    with pytest.raises(TypeError, match='dbt_task'):
+        jobs_service.submit_run(dbt_task=[])