provide an opportunity to parse db credentials from dbt profiles.yml

pik94 · pik94 · commit 24abee4e1afd · 2023-04-12T18:36:53.000+06:00
diff --git a/data_diff/cloud/data_source.py b/data_diff/cloud/data_source.py
@@ -1,4 +1,3 @@
-import enum
 import time
 from typing import List, Optional
 
@@ -14,6 +13,7 @@
     TDsConfig,
     TestDataSourceStatus,
 )
+from ..dbt_parser import DbtParser
 
 
 UNKNOWN_VALUE = "unknown_value"
@@ -49,8 +49,12 @@ def _validate_temp_schema(temp_schema: str):
         raise ValueError("Temporary schema should have a format <database>.<schema>")
 
 
-def create_ds_config(ds_config: TCloudApiDataSourceConfigSchema, data_source_name: str) -> TDsConfig:
-    options = _parse_ds_credentials(ds_config=ds_config, only_basic_settings=True)
+def create_ds_config(
+    ds_config: TCloudApiDataSourceConfigSchema,
+    data_source_name: str,
+    dbt_parser: Optional[DbtParser] = None,
+) -> TDsConfig:
+    options = _parse_ds_credentials(ds_config=ds_config, only_basic_settings=True, dbt_parser=dbt_parser)
 
     temp_schema = TemporarySchemaPrompt.ask("Temporary schema (<database>.<schema>)")
     float_tolerance = FloatPrompt.ask("Float tolerance", default=0.000001)
@@ -64,7 +68,25 @@ def create_ds_config(ds_config: TCloudApiDataSourceConfigSchema, data_source_nam
     )
 
 
-def _parse_ds_credentials(ds_config: TCloudApiDataSourceConfigSchema, only_basic_settings: bool = True):
+def _parse_ds_credentials(
+    ds_config: TCloudApiDataSourceConfigSchema, only_basic_settings: bool = True, dbt_parser: Optional[DbtParser] = None
+):
+    def cast_value(value: str, type_: str):
+        if type_ == "integer":
+            return int(value)
+        elif type_ == "boolean":
+            return bool(value)
+        return value
+
+    creds = {}
+    use_dbt_data = False
+    if dbt_parser is not None:
+        use_dbt_data = Confirm.ask("Would you like to extract database credentials from dbt profile.yml?")
+        try:
+            creds = dbt_parser.get_connection_creds()[0]
+        except Exception as e:
+            rich.print(f"[red]Cannot parse database credentials from dbt profile.yml. Reason: {e}")
+
     ds_options = {}
     basic_required_fields = set(ds_config.config_schema.required)
     for param_name, param_data in ds_config.config_schema.properties.items():
@@ -83,6 +105,14 @@ def _parse_ds_credentials(ds_config: TCloudApiDataSourceConfigSchema, only_basic
         if default_value != UNKNOWN_VALUE:
             input_values["default"] = default_value
 
+        if use_dbt_data:
+            value = creds.get(param_name, UNKNOWN_VALUE)
+            if value == UNKNOWN_VALUE:
+                rich.print(f'[red]Cannot extract "{param_name}" from dbt profile.yml. Please, type it manually')
+            else:
+                ds_options[param_name] = cast_value(value, type_)
+                continue
+
         if type_ == "integer":
             value = IntPrompt.ask(**input_values)
         elif type_ == "boolean":
@@ -177,7 +207,7 @@ def _render_data_source_test_results(test_results: List[TDataSourceTestStage]) -
     rich.print(table)
 
 
-def get_or_create_data_source(api: DatafoldAPI) -> int:
+def get_or_create_data_source(api: DatafoldAPI, dbt_parser: Optional[DbtParser] = None) -> int:
     ds_configs = api.get_data_source_schema_config()
     data_sources = api.get_data_sources()
 
@@ -198,10 +228,10 @@ def get_or_create_data_source(api: DatafoldAPI) -> int:
         _render_data_source(data_source=ds, title=f'Found existing data source for name "{ds.name}"')
         use_existing_ds = Confirm.ask("Would you like to continue with the existing data source?")
         if not use_existing_ds:
-            return get_or_create_data_source(api=api)
+            return get_or_create_data_source(api=api, dbt_parser=dbt_parser)
         return ds.id
 
-    ds_config = create_ds_config(ds_config, ds_name)
+    ds_config = create_ds_config(ds_config=ds_config, data_source_name=ds_name, dbt_parser=dbt_parser)
     ds = api.create_data_source(ds_config)
     data_source_url = f"{api.host}/settings/integrations/dwh/{ds.type}/{ds.id}"
     _render_data_source(data_source=ds, title=f"Created a new data source with ID = {ds.id} ({data_source_url})")
diff --git a/data_diff/dbt.py b/data_diff/dbt.py
@@ -1,21 +1,19 @@
-import json
 import os
 import time
 import webbrowser
 import rich
 from rich.prompt import Confirm
 
-from collections import defaultdict
 from dataclasses import dataclass
-from packaging.version import parse as parse_version
-from typing import List, Optional, Dict, Tuple, Set
+from typing import List, Optional, Dict
 from .utils import dbt_diff_string_template, getLogger
-from .version import __version__
 from pathlib import Path
 
 import keyring
 
 from .cloud import DatafoldAPI, TCloudApiDataDiff, get_or_create_data_source
+from .dbt_parser import DbtParser
+
 
 logger = getLogger(__name__)
 
@@ -41,7 +39,7 @@ def import_dbt():
     send_event_json,
     is_tracking_enabled,
 )
-from .utils import get_from_dict_with_raise, run_as_daemon, truncate_error
+from .utils import run_as_daemon, truncate_error
 from . import connect_to_table, diff_tables, Algorithm
 
 RUN_RESULTS_PATH = "target/run_results.json"
@@ -105,7 +103,7 @@ def dbt_diff(
             rich.print("[red]Data source ID not found in dbt_project.yml")
             is_create_data_source = Confirm.ask("Would you like to create a new data source?")
             if is_create_data_source:
-                datasource_id = get_or_create_data_source(api=api)
+                datasource_id = get_or_create_data_source(api=api, dbt_parser=dbt_parser)
                 rich.print(f'To use the data source in next runs, please, update your "{PROJECT_FILE}" with a block:')
                 rich.print(f"[green]vars:\n  data_diff:\n    datasource_id: {datasource_id}\n")
                 rich.print(
diff --git a/data_diff/dbt_parser.py b/data_diff/dbt_parser.py
