align dbt cred params with datafold cred params

Author: pik94

data_diff/cloud/data_source.py

@@ -1,3 +1,4 @@
+import json
 import time
 from typing import List, Optional, Union, overload
 
@@ -111,6 +112,29 @@ def _cast_value(value: str, type_: str) -> Union[bool, int, str]:
     return value
 
 
+def _get_data_from_bigquery_json(path: str):
+    with open(path, 'r') as file:
+        return json.load(file)
+
+
+def _align_dbt_cred_params_with_datafold_params(dbt_creds: dict) -> dict:
+    db_type = dbt_creds['type']
+    if db_type == 'bigquery':
+        method = dbt_creds['method']
+        if method == 'service-account':
+            data = _get_data_from_bigquery_json(path=dbt_creds['keyfile'])
+            dbt_creds['jsonKeyFile'] = json.dumps(data)
+        elif method == 'service-account-json':
+            dbt_creds['jsonKeyFile'] = json.dumps(dbt_creds['keyfile_json'])
+        dbt_creds['projectId'] = dbt_creds['project']
+    elif db_type == 'snowflake':
+        dbt_creds['default_db'] = dbt_creds['database']
+    elif db_type == 'databricks':
+        dbt_creds['http_password'] = dbt_creds['token']
+        dbt_creds['database'] = dbt_creds.get('catalog')
+    return dbt_creds
+
+
 def _parse_ds_credentials(
     ds_config: TCloudApiDataSourceConfigSchema, only_basic_settings: bool = True, dbt_parser: Optional[DbtParser] = None
 ):
@@ -120,6 +144,7 @@ def _parse_ds_credentials(
         use_dbt_data = Confirm.ask("Would you like to extract database credentials from dbt profiles.yml?")
         try:
             creds = dbt_parser.get_connection_creds()[0]
+            creds = _align_dbt_cred_params_with_datafold_params(dbt_creds=creds)
         except Exception as e:
             rich.print(f"[red]Cannot parse database credentials from dbt profiles.yml. Reason: {e}")
 

data_diff/cloud/datafold_api.py

@@ -1,3 +1,4 @@
+import base64
 import dataclasses
 import enum
 import time
@@ -159,7 +160,7 @@ class TCloudDataSourceTestResult(pydantic.BaseModel):
 class TCloudApiDataSourceTestResult(pydantic.BaseModel):
     name: str
     status: str
-    result: TCloudDataSourceTestResult
+    result: Optional[TCloudDataSourceTestResult]
 
 
 @dataclasses.dataclass
@@ -191,7 +192,11 @@ def get_data_sources(self) -> List[TCloudApiDataSource]:
         return [TCloudApiDataSource(**item) for item in rv.json()]
 
     def create_data_source(self, config: TDsConfig) -> TCloudApiDataSource:
-        rv = self.make_post_request(url="api/v1/data_sources", payload=config.dict())
+        payload = config.dict()
+        if config.type == 'bigquery':
+            json_string = payload['options']['jsonKeyFile'].encode('utf-8')
+            payload['options']['jsonKeyFile'] = base64.b64encode(json_string).decode('utf-8')
+        rv = self.make_post_request(url="api/v1/data_sources", payload=payload)
         return TCloudApiDataSource(**rv.json())
 
     def get_data_source_schema_config(
@@ -250,7 +255,7 @@ def check_data_source_test_results(self, job_id: int) -> List[TCloudApiDataSourc
                     status=item["result"]["code"].lower(),
                     message=item["result"]["message"],
                     outcome=item["result"]["outcome"],
-                ),
+                ) if item["result"] is not None else None,
             )
             for item in rv.json()["results"]
         ]