Merge branch 'master' into allow-dbt-selectors

Author: dlawin

data_diff/__init__.py

@@ -1,6 +1,6 @@
 from typing import Sequence, Tuple, Iterator, Optional, Union
 
-from sqeleton.abcs import DbTime, DbPath
+from data_diff.sqeleton.abcs import DbTime, DbPath
 
 from .tracking import disable_tracking
 from .databases import connect

data_diff/__main__.py

@@ -10,8 +10,8 @@
 import rich
 import click
 
-from sqeleton.schema import create_schema
-from sqeleton.queries.api import current_timestamp
+from data_diff.sqeleton.schema import create_schema
+from data_diff.sqeleton.queries.api import current_timestamp
 
 from .dbt import dbt_diff
 from .utils import eval_name_template, remove_password_from_url, safezip, match_like

data_diff/cloud/data_source.py

@@ -1,3 +1,4 @@
+import json
 import time
 from typing import List, Optional, Union, overload
 
@@ -50,14 +51,33 @@ def _validate_temp_schema(temp_schema: str):
         raise ValueError("Temporary schema should have a format <database>.<schema>")
 
 
+def _get_temp_schema(dbt_parser: DbtParser, db_type: str) -> Optional[str]:
+    diff_vars = dbt_parser.get_datadiff_variables()
+    config_prod_database = diff_vars.get("prod_database")
+    config_prod_schema = diff_vars.get("prod_schema")
+    if config_prod_database is not None and config_prod_schema is not None:
+        temp_schema = f"{config_prod_database}.{config_prod_schema}"
+        if db_type == "snowflake":
+            return temp_schema.upper()
+        elif db_type in {"pg", "postgres_aurora", "postgres_aws_rds", "redshift"}:
+            return temp_schema.lower()
+        return temp_schema
+    return
+
+
 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>)")
+    temp_schema = _get_temp_schema(dbt_parser=dbt_parser, db_type=ds_config.db_type) if dbt_parser else None
+    if temp_schema:
+        temp_schema = TemporarySchemaPrompt.ask("Temporary schema", default=temp_schema)
+    else:
+        temp_schema = TemporarySchemaPrompt.ask("Temporary schema (<database>.<schema>)")
+
     float_tolerance = FloatPrompt.ask("Float tolerance", default=0.000001)
 
     return TDsConfig(
@@ -92,6 +112,37 @@ 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"])
+        else:
+            rich.print(
+                f'[red]Cannot extract bigquery credentials from dbt_project.yml for "{method}" type. '
+                f"If you want to provide credentials via dbt_project.yml, "
+                f'please, use "service-account" or "service-account-json" '
+                f"(more in docs: https://docs.getdbt.com/reference/warehouse-setups/bigquery-setup). "
+                f"Otherwise, you can provide a path to a json key file or a json key file data as an input."
+            )
+        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
 ):
@@ -101,6 +152,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
@@ -6,6 +7,9 @@
 import pydantic
 import requests
 
+from ..utils import getLogger
+
+logger = getLogger(__name__)
 
 Self = TypeVar("Self", bound=pydantic.BaseModel)
 
@@ -97,6 +101,8 @@ class TCloudApiDataDiff(pydantic.BaseModel):
     table1: List[str]
     table2: List[str]
     pk_columns: List[str]
+    filter1: Optional[str] = None
+    filter2: Optional[str] = None
 
 
 class TSummaryResultPrimaryKeyStats(pydantic.BaseModel):
@@ -159,7 +165,7 @@ class TCloudDataSourceTestResult(pydantic.BaseModel):
 class TCloudApiDataSourceTestResult(pydantic.BaseModel):
     name: str
     status: str
-    result: TCloudDataSourceTestResult
+    result: Optional[TCloudDataSourceTestResult]
 
 
 @dataclasses.dataclass
@@ -191,7 +197,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(
@@ -216,11 +226,12 @@ def poll_data_diff_results(self, diff_id: int) -> TCloudApiDataDiffSummaryResult
         summary_results = None
         start_time = time.monotonic()
         sleep_interval = 5  # starts at 5 sec
-        max_sleep_interval = 60
+        max_sleep_interval = 30
         max_wait_time = 300
 
         diff_url = f"{self.host}/datadiffs/{diff_id}/overview"
         while not summary_results:
+            logger.debug(f"Polling: {diff_url}")
             response = self.make_get_request(url=f"api/v1/datadiffs/{diff_id}/summary_results")
             response_json = response.json()
             if response_json["status"] == "success":
@@ -250,7 +261,9 @@ 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"]
         ]

data_diff/databases/__init__.py

@@ -1,4 +1,4 @@
-from sqeleton.databases import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, QueryError, ConnectError
+from data_diff.sqeleton.databases import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, QueryError, ConnectError
 
 from .postgresql import PostgreSQL
 from .mysql import MySQL

data_diff/databases/_connect.py

@@ -1,6 +1,6 @@
 import logging
 
-from sqeleton.databases import Connect
+from data_diff.sqeleton.databases import Connect
 
 from .postgresql import PostgreSQL
 from .mysql import MySQL

data_diff/databases/base.py

@@ -1,4 +1,4 @@
-from sqeleton.abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue
+from data_diff.sqeleton.abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue
 
 
 class DatadiffDialect(AbstractMixin_MD5, AbstractMixin_NormalizeValue):

data_diff/databases/bigquery.py

@@ -1,4 +1,4 @@
-from sqeleton.databases import bigquery
+from data_diff.sqeleton.databases import bigquery
 from .base import DatadiffDialect
 
 

data_diff/databases/clickhouse.py

@@ -1,4 +1,4 @@
-from sqeleton.databases import clickhouse
+from data_diff.sqeleton.databases import clickhouse
 from .base import DatadiffDialect
 
 

data_diff/databases/databricks.py

@@ -1,4 +1,4 @@
-from sqeleton.databases import databricks
+from data_diff.sqeleton.databases import databricks
 from .base import DatadiffDialect