def test_get_run_with_dbt_cloud_bad_status_code(self, status_code, content):
self.setUp()
run_id = 1234
dbt_cloud_client = DbtCloudApiClient(
self.DBT_CLOUD_ACCOUNT_ID, self.DBT_CLOUD_API_KEY
)
server_error_response = Response()
server_error_response.status_code = status_code
server_error_response._content = content
session_mock = MagicMock()
session_mock.get.return_value = server_error_response
dbt_cloud_client.session = session_mock
res = dbt_cloud_client.get_run(run_id)
assert res is None
def test_get_run_with_connection_error_exception(self):
self.setUp()
run_id = 1234
dbt_cloud_client = DbtCloudApiClient(
self.DBT_CLOUD_ACCOUNT_ID, self.DBT_CLOUD_API_KEY
)
session_mock = MagicMock()
session_mock.side_effect = ConnectionError
dbt_cloud_client.session.get = session_mock
expected_endpoint = f"{dbt_cloud_client.administrative_api_url}{self.DBT_CLOUD_ACCOUNT_ID}/runs/{run_id}"
expected_params = {"include_related": '["run_steps"]'}
res = dbt_cloud_client.get_run(run_id)
session_mock.assert_called_with(url=expected_endpoint, params=expected_params)
assert res is None
def test_get_run_results_artifact(self):
self.setUp()
run_id = 1234
dbt_cloud_client = DbtCloudApiClient(
self.DBT_CLOUD_ACCOUNT_ID, self.DBT_CLOUD_API_KEY
)
dbt_cloud_client.send_request = self.send_request_mock
self.send_request_mock.return_value = {"run_id": run_id}
expected_endpoint = f"{dbt_cloud_client.administrative_api_url}{self.DBT_CLOUD_ACCOUNT_ID}/runs/{run_id}/artifacts/run_results.json"
expected_params = {"step": 1}
dbt_cloud_client.get_run_results_artifact(run_id)
self.send_request_mock.assert_called()
self.send_request_mock.assert_called_with(
endpoint=expected_endpoint, data=expected_params
)
def test_get_run(self):
self.setUp()
run_id = 1234
dbt_cloud_client = DbtCloudApiClient(
self.DBT_CLOUD_ACCOUNT_ID, self.DBT_CLOUD_API_KEY
)
dbt_cloud_client.send_request = self.send_request_mock
self.send_request_mock.return_value = {"run_id": run_id}
expected_endpoint = f"{dbt_cloud_client.administrative_api_url}{self.DBT_CLOUD_ACCOUNT_ID}/runs/{run_id}"
expected_params = {"include_related": '["run_steps"]'}
dbt_cloud_client.get_run(run_id)
self.send_request_mock.assert_called()
self.send_request_mock.assert_called_with(
endpoint=expected_endpoint, data=expected_params
)
def test_query_dbt_run_results(self):
self.setUp()
run_id = 1234
job_id = 5566
dbt_cloud_client = DbtCloudApiClient(
self.DBT_CLOUD_ACCOUNT_ID, self.DBT_CLOUD_API_KEY
)
dbt_cloud_client.send_request = self.send_request_mock
expected_query = "{\ntests(runId: %s,jobId: %s){\nuniqueId,\nstatus\n}\n}" % (
run_id,
job_id,
)
dbt_cloud_client.query_dbt_test_results(job_id, run_id)
self.send_request_mock.assert_called()
self.send_request_mock.assert_called_with(
endpoint=dbt_cloud_client.metadata_api_url,
method="POST",
data={"query": expected_query},
)
def collect_data_from_dbt_cloud(
dbt_cloud_account_id, dbt_cloud_api_token, dbt_job_run_id
):
"""
Collect metadata for a single run from dbt cloud.
Args:
dbt_cloud_account_id: dbt cloud account id in order to identify with dbt cloud API
dbt_cloud_api_token: Api token in order to authenticate dbt cloud API
dbt_job_run_id: run id of the dbt run that we want to report it's metadata.
@task
def prepare_data():
collect_data_from_dbt_cloud(
dbt_cloud_account_id=my_dbt_cloud_account_id,
dbt_cloud_api_token="my_dbt_cloud_api_token",
dbt_job_run_id=12345
)
"""
if not dbt_cloud_api_token or not dbt_cloud_account_id:
logger.warning(
"Can't collect Data from dbt cloud, account id nor api key were supplied"
)
return
try:
dbt_cloud_client = DbtCloudApiClient(
account_id=dbt_cloud_account_id, dbt_cloud_api_token=dbt_cloud_api_token
)
dbt_run_meta_data = dbt_cloud_client.get_run(dbt_job_run_id)
if not dbt_run_meta_data:
logger.warning("Fail getting run data from dbt cloud ")
return
env_id = dbt_run_meta_data.get("environment_id")
env = dbt_cloud_client.get_environment(env_id=env_id)
if env:
dbt_run_meta_data["environment"] = env
for step in dbt_run_meta_data.get("run_steps", []):
step_run_results_artifact = dbt_cloud_client.get_run_results_artifact(
dbt_job_run_id, step["index"]
)
if step_run_results_artifact:
step["run_results"] = step_run_results_artifact
step_run_manifest_artifact = dbt_cloud_client.get_manifest_artifact(
dbt_job_run_id, step["index"]
)
if step_run_manifest_artifact:
step["manifest"] = step_run_manifest_artifact
_report_dbt_metadata(dbt_run_meta_data)
except Exception as e:
log_exception("Could not collect data from dbt cloud", e)