override --cloud host name

Author: dlawin

data_diff/__main__.py

@@ -228,6 +228,12 @@ def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -
     metavar="PATH",
     help="Which directory to look in for the dbt_project.yml file. Default is the current working directory and its parents.",
 )
+@click.option(
+    "--cloud-host-name",
+    envvar="DATAFOLD_API_HOST_NAME",
+    default=None,
+    help="Override the host name for on-premise Datafold deployments --cloud api calls. Can also be set via the DATAFOLD_API_HOST_NAME environment variable.",
+)
 def main(conf, run, **kw):
     if kw["table2"] is None and kw["database2"]:
         # Use the "database table table" form
@@ -264,6 +270,7 @@ def main(conf, run, **kw):
                 profiles_dir_override=kw["dbt_profiles_dir"],
                 project_dir_override=kw["dbt_project_dir"],
                 is_cloud=kw["cloud"],
+                cloud_host_name=kw["cloud_host_name"],
             )
         else:
             return _data_diff(**kw)
@@ -306,6 +313,7 @@ def _data_diff(
     cloud,
     dbt_profiles_dir,
     dbt_project_dir,
+    cloud_host_name,
     threads1=None,
     threads2=None,
     __conf__=None,

data_diff/dbt.py

@@ -73,7 +73,10 @@ class DiffVars:
 
 
 def dbt_diff(
-    profiles_dir_override: Optional[str] = None, project_dir_override: Optional[str] = None, is_cloud: bool = False
+    profiles_dir_override: Optional[str] = None,
+    project_dir_override: Optional[str] = None,
+    is_cloud: bool = False,
+    cloud_host_name: Optional[str] = None,
 ) -> None:
     set_entrypoint_name("CLI-dbt")
     dbt_parser = DbtParser(profiles_dir_override, project_dir_override, is_cloud)
@@ -101,7 +104,7 @@ def dbt_diff(
         )
 
         if is_cloud and len(diff_vars.primary_keys) > 0:
-            _cloud_diff(diff_vars)
+            _cloud_diff(diff_vars, cloud_host_name)
         elif not is_cloud and len(diff_vars.primary_keys) > 0:
             _local_diff(diff_vars)
         else:
@@ -219,8 +222,9 @@ def _local_diff(diff_vars: DiffVars) -> None:
         )
 
 
-def _cloud_diff(diff_vars: DiffVars) -> None:
+def _cloud_diff(diff_vars: DiffVars, host_name: Optional[str] = None) -> None:
     api_key = os.environ.get("DATAFOLD_API_KEY")
+    host_name = host_name or "app.datafold.com"
 
     if diff_vars.datasource_id is None:
         raise ValueError(
@@ -229,7 +233,7 @@ def _cloud_diff(diff_vars: DiffVars) -> None:
     if api_key is None:
         raise ValueError("API key not found, add it as an environment variable called DATAFOLD_API_KEY.")
 
-    url = "https://app.datafold.com/api/v1/datadiffs"
+    url = f"https://{host_name}/api/v1/datadiffs"
 
     payload = {
         "data_source1_id": diff_vars.datasource_id,
@@ -255,8 +259,8 @@ def _cloud_diff(diff_vars: DiffVars) -> None:
         response.raise_for_status()
         data = response.json()
         diff_id = data["id"]
-        # TODO in future we should support self hosted datafold
-        diff_url = f"https://app.datafold.com/datadiffs/{diff_id}/overview"
+
+        diff_url = f"https://{host_name}/datadiffs/{diff_id}/overview"
         rich.print(
             "[red]"
             + ".".join(diff_vars.prod_path)

tests/test_dbt.py

@@ -426,11 +426,48 @@ def test_cloud_diff(self, mock_request, mock_os_environ, mock_print):
 
         mock_request.assert_called_once()
         mock_print.assert_called_once()
+        request_endpoint = mock_request.call_args[0][1]
         request_data_dict = mock_request.call_args[1]["json"]
         self.assertEqual(
             mock_request.call_args[1]["headers"]["Authorization"],
             "Key " + expected_api_key,
         )
+        self.assertEqual(request_endpoint, f'https://app.datafold.com/api/v1/datadiffs')
+        self.assertEqual(request_data_dict["data_source1_id"], expected_datasource_id)
+        self.assertEqual(request_data_dict["data_source2_id"], expected_datasource_id)
+        self.assertEqual(request_data_dict["table1"], prod_qualified_list)
+        self.assertEqual(request_data_dict["table2"], dev_qualified_list)
+        self.assertEqual(request_data_dict["pk_columns"], expected_primary_keys)
+        
+    @patch("data_diff.dbt.rich.print")
+    @patch("data_diff.dbt.os.environ")
+    @patch("data_diff.dbt.requests.request")
+    def test_cloud_diff_host_name_override(self, mock_request, mock_os_environ, mock_print):
+        expected_api_key = "an_api_key"
+        mock_response = Mock()
+        mock_response.json.return_value = {"id": 123}
+        mock_request.return_value = mock_response
+        mock_os_environ.get.return_value = expected_api_key
+        dev_qualified_list = ["dev_db", "dev_schema", "dev_table"]
+        prod_qualified_list = ["prod_db", "prod_schema", "prod_table"]
+        expected_datasource_id = 1
+        expected_primary_keys = ["primary_key_column"]
+        diff_vars = DiffVars(
+            dev_qualified_list, prod_qualified_list, expected_primary_keys, expected_datasource_id, None, None
+        )
+        host_name = "a_host_name"
+        _cloud_diff(diff_vars, host_name)
+
+        mock_request.assert_called_once()
+        mock_print.assert_called_once()
+        
+        request_endpoint = mock_request.call_args[0][1]
+        request_data_dict = mock_request.call_args[1]["json"]
+        self.assertEqual(
+            mock_request.call_args[1]["headers"]["Authorization"],
+            "Key " + expected_api_key,
+        )
+        self.assertEqual(request_endpoint, f'https://{host_name}/api/v1/datadiffs')
         self.assertEqual(request_data_dict["data_source1_id"], expected_datasource_id)
         self.assertEqual(request_data_dict["data_source2_id"], expected_datasource_id)
         self.assertEqual(request_data_dict["table1"], prod_qualified_list)
@@ -500,7 +537,35 @@ def test_diff_is_cloud(self, mock_print, mock_dbt_parser, mock_cloud_diff, mock_
         mock_dbt_parser_inst.get_models.assert_called_once()
         mock_dbt_parser_inst.set_connection.assert_not_called()
 
-        mock_cloud_diff.assert_called_once_with(expected_diff_vars)
+        mock_cloud_diff.assert_called_once_with(expected_diff_vars, None)
+        mock_local_diff.assert_not_called()
+        mock_print.assert_called_once()
+        
+    @patch("data_diff.dbt._get_diff_vars")
+    @patch("data_diff.dbt._local_diff")
+    @patch("data_diff.dbt._cloud_diff")
+    @patch("data_diff.dbt.DbtParser.__new__")
+    @patch("data_diff.dbt.rich.print")
+    def test_diff_is_cloud(self, mock_print, mock_dbt_parser, mock_cloud_diff, mock_local_diff, mock_get_diff_vars):
+        mock_dbt_parser_inst = Mock()
+        mock_model = Mock()
+        expected_dbt_vars_dict = {
+            "prod_database": "prod_db",
+            "prod_schema": "prod_schema",
+            "datasource_id": 1,
+        }
+        host_name = 'a_host_name'
+
+        mock_dbt_parser.return_value = mock_dbt_parser_inst
+        mock_dbt_parser_inst.get_models.return_value = [mock_model]
+        mock_dbt_parser_inst.get_datadiff_variables.return_value = expected_dbt_vars_dict
+        expected_diff_vars = DiffVars(["dev"], ["prod"], ["pks"], 123, None, None)
+        mock_get_diff_vars.return_value = expected_diff_vars
+        dbt_diff(is_cloud=True, cloud_host_name=host_name)
+        mock_dbt_parser_inst.get_models.assert_called_once()
+        mock_dbt_parser_inst.set_connection.assert_not_called()
+
+        mock_cloud_diff.assert_called_once_with(expected_diff_vars, host_name)
         mock_local_diff.assert_not_called()
         mock_print.assert_called_once()