Merge pull request #421 from datafold/optional_dbt

Make DBT dependency optional

Author: erezsh

data_diff/dbt.py

@@ -3,14 +3,21 @@
 import os
 import time
 import rich
-import yaml
 from dataclasses import dataclass
 from packaging.version import parse as parse_version
 from typing import List, Optional, Dict
 
 import requests
-from dbt_artifacts_parser.parser import parse_run_results, parse_manifest
-from dbt.config.renderer import ProfileRenderer
+
+def import_dbt():
+ try:
+ from dbt_artifacts_parser.parser import parse_run_results, parse_manifest
+ from dbt.config.renderer import ProfileRenderer
+ import yaml
+ except ImportError:
+ raise RuntimeError("Could not import 'dbt' package. You can install it using: pip install 'data-diff[dbt]'.")
+
+ return parse_run_results, parse_manifest, ProfileRenderer, yaml
 
 from .tracking import (
  set_entrypoint_name,
@@ -263,13 +270,15 @@ def __init__(self, profiles_dir_override: str, project_dir_override: str, is_clo
  self.project_dict = None
  self.requires_upper = False
 
+ self.parse_run_results, self.parse_manifest, self.ProfileRenderer, self.yaml = import_dbt()
+
  def get_datadiff_variables(self) -> dict:
  return self.project_dict.get("vars").get("data_diff")
 
  def get_models(self):
  with open(self.project_dir + RUN_RESULTS_PATH) as run_results:
  run_results_dict = json.load(run_results)
- run_results_obj = parse_run_results(run_results=run_results_dict)
+ run_results_obj = self.parse_run_results(run_results=run_results_dict)
 
  dbt_version = parse_version(run_results_obj.metadata.dbt_version)
 
@@ -280,7 +289,7 @@ def get_models(self):
 
  with open(self.project_dir + MANIFEST_PATH) as manifest:
  manifest_dict = json.load(manifest)
- manifest_obj = parse_manifest(manifest=manifest_dict)
+ manifest_obj = self.parse_manifest(manifest=manifest_dict)
 
  success_models = [x.unique_id for x in run_results_obj.results if x.status.name == "success"]
  models = [manifest_obj.nodes.get(x) for x in success_models]
@@ -295,11 +304,11 @@ def get_primary_keys(self, model):
 
  def set_project_dict(self):
  with open(self.project_dir + PROJECT_FILE) as project:
- self.project_dict = yaml.safe_load(project)
+ self.project_dict = self.yaml.safe_load(project)
 
  def set_connection(self):
  with open(self.profiles_dir + PROFILES_FILE) as profiles:
- profiles = yaml.safe_load(profiles)
+ profiles = self.yaml.safe_load(profiles)
 
  dbt_profile = self.project_dict.get("profile")
  profile_outputs = profiles.get(dbt_profile)
@@ -308,7 +317,7 @@ def set_connection(self):
  conn_type = credentials.get("type").lower()
 
  # values can contain env_vars
- rendered_credentials = ProfileRenderer().render_data(credentials)
+ rendered_credentials = self.ProfileRenderer().render_data(credentials)
 
  if conn_type == "snowflake":
  if rendered_credentials.get("password") is None or rendered_credentials.get("private_key_path") is not None:

pyproject.toml

@@ -38,8 +38,8 @@ trino = {version="^0.314.0", optional=true}
 presto-python-client = {version="*", optional=true}
 clickhouse-driver = {version="*", optional=true}
 duckdb = {version="^0.6.0", optional=true}
-dbt-artifacts-parser = "^0.2.4"
-dbt-core = "^1.0.0"
+dbt-artifacts-parser = {version="^0.2.4", optional=true}
+dbt-core = {version="^1.0.0", optional=true}
 
 [tool.poetry.dev-dependencies]
 parameterized = "*"
@@ -54,6 +54,8 @@ presto-python-client = "*"
 clickhouse-driver = "*"
 vertica-python = "*"
 duckdb = "^0.6.0"
+dbt-artifacts-parser = "^0.2.4"
+dbt-core = "^1.0.0"
 # google-cloud-bigquery = "*"
 # databricks-sql-connector = "*"
 
@@ -70,6 +72,7 @@ trino = ["trino"]
 clickhouse = ["clickhouse-driver"]
 vertica = ["vertica-python"]
 duckdb = ["duckdb"]
+dbt = ["dbt-core", "dbt-artifacts-parser"]
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

tests/test_dbt.py

@@ -1,6 +1,7 @@
 import json
 import os
 
+import yaml
 from data_diff.diff_tables import Algorithm
 from .test_cli import run_datadiff_cli
 
@@ -49,112 +50,102 @@ def test_get_datadiff_variables_empty(self):
  DbtParser.get_datadiff_variables(mock_self)
 
  @patch("builtins.open", new_callable=mock_open, read_data="{}")
- @patch("data_diff.dbt.parse_run_results")
- @patch("data_diff.dbt.parse_manifest")
- def test_get_models(self, mock_manifest_parser, mock_run_parser, mock_open):
+ def test_get_models(self, mock_open):
  expected_value = "expected_value"
  mock_self = Mock()
  mock_self.project_dir = ""
  mock_run_results = Mock()
  mock_success_result = Mock()
  mock_failed_result = Mock()
  mock_manifest = Mock()
- mock_run_parser.return_value = mock_run_results
+ mock_self.parse_run_results.return_value = mock_run_results
  mock_run_results.metadata.dbt_version = "1.0.0"
  mock_success_result.unique_id = "success_unique_id"
  mock_failed_result.unique_id = "failed_unique_id"
  mock_success_result.status.name = "success"
  mock_failed_result.status.name = "failed"
  mock_run_results.results = [mock_success_result, mock_failed_result]
- mock_manifest_parser.return_value = mock_manifest
+ mock_self.parse_manifest.return_value = mock_manifest
  mock_manifest.nodes = {"success_unique_id": expected_value}
 
  models = DbtParser.get_models(mock_self)
 
  self.assertEqual(expected_value, models[0])
  mock_open.assert_any_call(RUN_RESULTS_PATH)
  mock_open.assert_any_call(MANIFEST_PATH)
- mock_run_parser.assert_called_once_with(run_results={})
- mock_manifest_parser.assert_called_once_with(manifest={})
+ mock_self.parse_run_results.assert_called_once_with(run_results={})
+ mock_self.parse_manifest.assert_called_once_with(manifest={})
 
  @patch("builtins.open", new_callable=mock_open, read_data="{}")
- @patch("data_diff.dbt.parse_run_results")
- @patch("data_diff.dbt.parse_manifest")
- def test_get_models_bad_lower_dbt_version(self, mock_manifest_parser, mock_run_parser, mock_open):
+ def test_get_models_bad_lower_dbt_version(self, mock_open):
  mock_self = Mock()
  mock_self.project_dir = ""
  mock_run_results = Mock()
- mock_run_parser.return_value = mock_run_results
+ mock_self.parse_run_results.return_value = mock_run_results
  mock_run_results.metadata.dbt_version = "0.19.0"
 
  with self.assertRaises(Exception) as ex:
  DbtParser.get_models(mock_self)
 
  mock_open.assert_called_once_with(RUN_RESULTS_PATH)
- mock_run_parser.assert_called_once_with(run_results={})
- mock_manifest_parser.assert_not_called()
+ mock_self.parse_run_results.assert_called_once_with(run_results={})
+ mock_self.parse_manifest.assert_not_called()
  self.assertIn("version to be", ex.exception.args[0])
 
  @patch("builtins.open", new_callable=mock_open, read_data="{}")
- @patch("data_diff.dbt.parse_run_results")
- @patch("data_diff.dbt.parse_manifest")
- def test_get_models_bad_upper_dbt_version(self, mock_manifest_parser, mock_run_parser, mock_open):
+ def test_get_models_bad_upper_dbt_version(self, mock_open):
  mock_self = Mock()
  mock_self.project_dir = ""
  mock_run_results = Mock()
- mock_run_parser.return_value = mock_run_results
+ mock_self.parse_run_results.return_value = mock_run_results
  mock_run_results.metadata.dbt_version = "1.5.1"
 
  with self.assertRaises(Exception) as ex:
  DbtParser.get_models(mock_self)
 
  mock_open.assert_called_once_with(RUN_RESULTS_PATH)
- mock_run_parser.assert_called_once_with(run_results={})
- mock_manifest_parser.assert_not_called()
+ mock_self.parse_run_results.assert_called_once_with(run_results={})
+ mock_self.parse_manifest.assert_not_called()
  self.assertIn("version to be", ex.exception.args[0])
 
  @patch("builtins.open", new_callable=mock_open, read_data="{}")
- @patch("data_diff.dbt.parse_run_results")
- @patch("data_diff.dbt.parse_manifest")
- def test_get_models_no_success(self, mock_manifest_parser, mock_run_parser, mock_open):
+ def test_get_models_no_success(self, mock_open):
  mock_self = Mock()
  mock_self.project_dir = ""
  mock_run_results = Mock()
  mock_success_result = Mock()
  mock_failed_result = Mock()
  mock_manifest = Mock()
- mock_run_parser.return_value = mock_run_results
+ mock_self.parse_run_results.return_value = mock_run_results
  mock_run_results.metadata.dbt_version = "1.0.0"
  mock_failed_result.unique_id = "failed_unique_id"
  mock_success_result.status.name = "success"
  mock_failed_result.status.name = "failed"
  mock_run_results.results = [mock_failed_result]
- mock_manifest_parser.return_value = mock_manifest
+ mock_self.parse_manifest.return_value = mock_manifest
  mock_manifest.nodes = {"success_unique_id": "a_unique_id"}
 
  with self.assertRaises(Exception):
  DbtParser.get_models(mock_self)
 
  mock_open.assert_any_call(RUN_RESULTS_PATH)
  mock_open.assert_any_call(MANIFEST_PATH)
- mock_run_parser.assert_called_once_with(run_results={})
- mock_manifest_parser.assert_called_once_with(manifest={})
+ mock_self.parse_run_results.assert_called_once_with(run_results={})
+ mock_self.parse_manifest.assert_called_once_with(manifest={})
 
- @patch("data_diff.dbt.yaml.safe_load")
  @patch("builtins.open", new_callable=mock_open, read_data="key:\n value")
- def test_set_project_dict(self, mock_open, mock_yaml_parse):
+ def test_set_project_dict(self, mock_open):
  expected_dict = {"key1": "value1"}
  mock_self = Mock()
  mock_self.project_dir = ""
- mock_yaml_parse.return_value = expected_dict
+ mock_self.yaml.safe_load.return_value = expected_dict
  DbtParser.set_project_dict(mock_self)
 
  self.assertEqual(mock_self.project_dict, expected_dict)
  mock_open.assert_called_once_with(PROJECT_FILE)
 
- @patch("data_diff.dbt.yaml.safe_load")
  @patch("builtins.open", new_callable=mock_open, read_data="key:\n value")
- def test_set_connection_snowflake(self, mock_open_file, mock_yaml_parse):
+ def test_set_connection_snowflake(self, mock_open_file):
  expected_driver = "snowflake"
  expected_password = "password_value"
  profiles_dict = {
@@ -172,19 +163,19 @@ def test_set_connection_snowflake(self, mock_open_file, mock_yaml_parse):
  mock_self = Mock()
  mock_self.profiles_dir = ""
  mock_self.project_dict = {"profile": "profile_name"}
- mock_yaml_parse.return_value = profiles_dict
+ mock_self.yaml.safe_load.return_value = profiles_dict
+ mock_self.ProfileRenderer().render_data.return_value = profiles_dict["profile_name"]["outputs"]["connection1"]
  DbtParser.set_connection(mock_self)
 
  self.assertIsInstance(mock_self.connection, dict)
  self.assertEqual(mock_self.connection.get("driver"), expected_driver)
  self.assertEqual(mock_self.connection.get("password"), expected_password)
  self.assertEqual(mock_self.requires_upper, True)
  mock_open_file.assert_called_once_with(PROFILES_FILE)
- mock_yaml_parse.assert_called_once_with(mock_open_file())
+ mock_self.yaml.safe_load.assert_called_once_with(mock_open_file())
 
- @patch("data_diff.dbt.yaml.safe_load")
  @patch("builtins.open", new_callable=mock_open, read_data="key:\n value")
- def test_set_connection_snowflake_no_password(self, mock_open_file, mock_yaml_parse):
+ def test_set_connection_snowflake_no_password(self, mock_open_file):
  expected_driver = "snowflake"
  profiles_dict = {
  "profile_name": {
@@ -196,18 +187,18 @@ def test_set_connection_snowflake_no_password(self, mock_open_file, mock_yaml_pa
  mock_self = Mock()
  mock_self.profiles_dir = ""
  mock_self.project_dict = {"profile": "profile_name"}
- mock_yaml_parse.return_value = profiles_dict
+ mock_self.yaml.safe_load.return_value = profiles_dict
+ mock_self.ProfileRenderer().render_data.return_value = profiles_dict["profile_name"]["outputs"]["connection1"]
 
  with self.assertRaises(Exception):
  DbtParser.set_connection(mock_self)
 
  mock_open_file.assert_called_once_with(PROFILES_FILE)
- mock_yaml_parse.assert_called_once_with(mock_open_file())
+ mock_self.yaml.safe_load.assert_called_once_with(mock_open_file())
  self.assertNotIsInstance(mock_self.connection, dict)
 
- @patch("data_diff.dbt.yaml.safe_load")
  @patch("builtins.open", new_callable=mock_open, read_data="key:\n value")
- def test_set_connection_bigquery(self, mock_open_file, mock_yaml_parse):
+ def test_set_connection_bigquery(self, mock_open_file):
  expected_driver = "bigquery"
  expected_method = "oauth"
  expected_project = "a_project"
@@ -229,19 +220,19 @@ def test_set_connection_bigquery(self, mock_open_file, mock_yaml_parse):
  mock_self = Mock()
  mock_self.profiles_dir = ""
  mock_self.project_dict = {"profile": "profile_name"}
- mock_yaml_parse.return_value = profiles_dict
+ mock_self.yaml.safe_load.return_value = profiles_dict
+ mock_self.ProfileRenderer().render_data.return_value = profiles_dict["profile_name"]["outputs"]["connection1"]
  DbtParser.set_connection(mock_self)
 
  self.assertIsInstance(mock_self.connection, dict)
  self.assertEqual(mock_self.connection.get("driver"), expected_driver)
  self.assertEqual(mock_self.connection.get("project"), expected_project)
  self.assertEqual(mock_self.connection.get("dataset"), expected_dataset)
  mock_open_file.assert_called_once_with(PROFILES_FILE)
- mock_yaml_parse.assert_called_once_with(mock_open_file())
+ mock_self.yaml.safe_load.assert_called_once_with(mock_open_file())
 
- @patch("data_diff.dbt.yaml.safe_load")
  @patch("builtins.open", new_callable=mock_open, read_data="key:\n value")
- def test_set_connection_bigquery_not_oauth(self, mock_open_file, mock_yaml_parse):
+ def test_set_connection_bigquery_not_oauth(self, mock_open_file):
  expected_driver = "bigquery"
  expected_method = "not_oauth"
  expected_project = "a_project"
@@ -263,17 +254,17 @@ def test_set_connection_bigquery_not_oauth(self, mock_open_file, mock_yaml_parse
  mock_self = Mock()
  mock_self.profiles_dir = ""
  mock_self.project_dict = {"profile": "profile_name"}
- mock_yaml_parse.return_value = profiles_dict
+ mock_self.yaml.safe_load.return_value = profiles_dict
+ mock_self.ProfileRenderer().render_data.return_value = profiles_dict["profile_name"]["outputs"]["connection1"]
  with self.assertRaises(Exception):
  DbtParser.set_connection(mock_self)
 
  mock_open_file.assert_called_once_with(PROFILES_FILE)
- mock_yaml_parse.assert_called_once_with(mock_open_file())
+ mock_self.yaml.safe_load.assert_called_once_with(mock_open_file())
  self.assertNotIsInstance(mock_self.connection, dict)
 
- @patch("data_diff.dbt.yaml.safe_load")
  @patch("builtins.open", new_callable=mock_open, read_data="key:\n value")
- def test_set_connection_key_error(self, mock_open_file, mock_yaml_parse):
+ def test_set_connection_key_error(self, mock_open_file):
  profiles_dict = {
  "profile_name": {
  "outputs": {
@@ -290,17 +281,17 @@ def test_set_connection_key_error(self, mock_open_file, mock_yaml_parse):
  mock_self.profiles_dir = ""
  mock_self.project_dir = ""
  mock_self.project_dict = {"profile": "bad_key"}
- mock_yaml_parse.return_value = profiles_dict
+ mock_self.yaml.safe_load.return_value = profiles_dict
+ mock_self.ProfileRenderer().render_data.return_value = profiles_dict["profile_name"]["outputs"]["connection1"]
  with self.assertRaises(Exception):
  DbtParser.set_connection(mock_self)
 
  mock_open_file.assert_called_once_with(PROFILES_FILE)
- mock_yaml_parse.assert_called_once_with(mock_open_file())
+ mock_self.yaml.safe_load.assert_called_once_with(mock_open_file())
  self.assertNotIsInstance(mock_self.connection, dict)
 
- @patch("data_diff.dbt.yaml.safe_load")
  @patch("builtins.open", new_callable=mock_open, read_data="key:\n value")
- def test_set_connection_not_implemented(self, mock_open_file, mock_yaml_parse):
+ def test_set_connection_not_implemented(self, mock_open_file):
  expected_driver = "not_implemented"
  profiles_dict = {
  "profile_name": {
@@ -317,12 +308,13 @@ def test_set_connection_not_implemented(self, mock_open_file, mock_yaml_parse):
  mock_self.profiles_dir = ""
  mock_self.project_dir = ""
  mock_self.project_dict = {"profile": "profile_name"}
- mock_yaml_parse.return_value = profiles_dict
+ mock_self.yaml.safe_load.return_value = profiles_dict
+ mock_self.ProfileRenderer().render_data.return_value = profiles_dict["profile_name"]["outputs"]["connection1"]
  with self.assertRaises(NotImplementedError):
  DbtParser.set_connection(mock_self)
 
  mock_open_file.assert_called_once_with(PROFILES_FILE)
- mock_yaml_parse.assert_called_once_with(mock_open_file())
+ mock_self.yaml.safe_load.assert_called_once_with(mock_open_file())
  self.assertNotIsInstance(mock_self.connection, dict)