Make DuckDB data diffs work better (#716)

* duckdb joindiff works!

* fix example

* add motherduck example

* show a toml example with motherduck

* show that it's an example

* remove sensitive token from logs

* add token tests

* scrub token passing tests

* add duckdb to joindiff tests

---------

Co-authored-by: Sung Won Chung <[email protected]>

Author: sungchun12

README.md

@@ -81,18 +81,73 @@ More information about the algorithm and performance considerations can be found
 pip install data-diff 'data-diff[postgresql,snowflake]' -U
 ```
 
-Run `data-diff` with connection URIs. In the following example, we compare tables between PostgreSQL and Snowflake using hashdiff algorithm:
-```
+Run `data-diff` with connection URIs. In the following example, we compare tables between PostgreSQL and Snowflake using the hashdiff algorithm:
+
+```bash
 data-diff \
   postgresql://<username>:'<password>'@localhost:5432/<database> \
   <table> \
-  "snowflake://<username>:<password>@<password>/<DATABASE>/<SCHEMA>?warehouse=<WAREHOUSE>&role=<ROLE>" \
+  "snowflake://<username>:<password>@<account>/<DATABASE>/<SCHEMA>?warehouse=<WAREHOUSE>&role=<ROLE>" \
   <TABLE> \
   -k <primary key column> \
   -c <columns to compare> \
   -w <filter condition>
 ```
 
+Run `data-diff` with a `toml` configuration file. In the following example, we compare tables between MotherDuck(hosted DuckDB) and Snowflake using the hashdiff algorithm:
+
+```toml
+## DATABASE CONNECTION ##
+[database.duckdb_connection] 
+  driver = "duckdb"
+  # filepath = "datafold_demo.duckdb" # local duckdb file example
+  # filepath = "md:" # default motherduck connection example
+  filepath = "md:datafold_demo?motherduck_token=${motherduck_token}" # API token recommended for motherduck connection
+  database = "datafold_demo"
+
+[database.snowflake_connection]
+  driver = "snowflake"
+  database = "DEV"
+  user = "sung"
+  password = "${SNOWFLAKE_PASSWORD}" # or "<PASSWORD_STRING>"
+  # the info below is only required for snowflake
+  account = "${ACCOUNT}" # by33919
+  schema = "DEVELOPMENT"
+  warehouse = "DEMO"
+  role = "DEMO_ROLE"
+
+## RUN PARAMETERS ##
+[run.default]
+  verbose = true
+
+## EXAMPLE DATA DIFF JOB ##
+[run.demo_xdb_diff]
+  # Source 1 ("left")
+  1.database = "duckdb_connection"
+  1.table = "development.raw_orders"
+
+  # Source 2 ("right")
+  2.database = "snowflake_connection"
+  2.table = "RAW_ORDERS" # note that snowflake table names are case-sensitive
+
+  verbose = false
+```
+
+```bash
+# export relevant environment variables, example below
+export motherduck_token=<MOTHERDUCK_TOKEN>
+
+# run the configured data-diff job
+data-diff --conf datadiff.toml \
+  --run demo_xdb_diff \
+  -k "id" \
+  -c status
+
+# output example
+- 1, completed
++ 1, returned
+```
+
 Check out [documentation](https://docs.datafold.com/reference/open_source/cli) for the full command reference.
 
 
@@ -106,13 +161,14 @@ Check out [documentation](https://docs.datafold.com/reference/open_source/cli) f
 | Snowflake     |  🟢     | `"snowflake://<user>[:<password>]@<account>/<database>/<SCHEMA>?warehouse=<WAREHOUSE>&role=<role>[&authenticator=externalbrowser]"` |
 | BigQuery      |  🟢     | `bigquery://<project>/<dataset>`                                                                                                    |
 | Redshift      |  🟢     | `redshift://<username>:<password>@<hostname>:5439/<database>`                                                                       |
+| DuckDB        |  🟢   | `duckdb://<dbname>@<filepath>`                                                                                          |
+| MotherDuck        |  🟢   | `duckdb://<dbname>@<filepath>`                                                                                                   |
 | Oracle        |  🟡   | `oracle://<username>:<password>@<hostname>/servive_or_sid`                                                                          |
 | Presto        |  🟡   | `presto://<username>:<password>@<hostname>:8080/<database>`                                                                         |
 | Databricks    |  🟡   | `databricks://<http_path>:<access_token>@<server_hostname>/<catalog>/<schema>`                                                      |
 | Trino         |  🟡   | `trino://<username>:<password>@<hostname>:8080/<database>`                                                                          |
 | Clickhouse    |  🟡   | `clickhouse://<username>:<password>@<hostname>:9000/<database>`                                                                     |
 | Vertica       |  🟡   | `vertica://<username>:<password>@<hostname>:5433/<database>`                                                                        |
-| DuckDB        |  🟡   |                                                                                                                                     |
 | ElasticSearch |  📝    |                                                                                                                                     |
 | Planetscale   |  📝    |                                                                                                                                     |
 | Pinot         |  📝    |                                                                                                                                     |

data_diff/__main__.py

@@ -59,10 +59,13 @@ def _get_log_handlers(is_dbt: Optional[bool] = False) -> Dict[str, logging.Handl
     return handlers
 
 
-def _remove_passwords_in_dict(d: dict):
+def _remove_passwords_in_dict(d: dict) -> None:
     for k, v in d.items():
         if k == "password":
             d[k] = "*" * len(v)
+        elif k == "filepath":
+            if "motherduck_token=" in v:
+                d[k] = v.split("motherduck_token=")[0] + "motherduck_token=**********"
         elif isinstance(v, dict):
             _remove_passwords_in_dict(v)
         elif k.startswith("database"):

data_diff/joindiff_tables.py

@@ -9,7 +9,7 @@
 
 import attrs
 
-from data_diff.databases import Database, MsSQL, MySQL, BigQuery, Presto, Oracle, Snowflake
+from data_diff.databases import Database, MsSQL, MySQL, BigQuery, Presto, Oracle, Snowflake, DuckDB
 from data_diff.abcs.database_types import NumericType, DbPath
 from data_diff.databases.base import Compiler
 from data_diff.queries.api import (
@@ -157,7 +157,7 @@ def _diff_tables_root(self, table1: TableSegment, table2: TableSegment, info_tre
             drop_table(db, self.materialize_to_table)
 
         with self._run_in_background(*bg_funcs):
-            if isinstance(db, (Snowflake, BigQuery)):
+            if isinstance(db, (Snowflake, BigQuery, DuckDB)):
                 # Don't segment the table; let the database handling parallelization
                 yield from self._diff_segments(None, table1, table2, info_tree, None)
             else:

data_diff/utils.py

@@ -270,6 +270,9 @@ def remove_passwords_in_dict(d: dict, replace_with: str = "***"):
     for k, v in d.items():
         if k == "password":
             d[k] = replace_with
+        elif k == "filepath":
+            if "motherduck_token=" in v:
+                d[k] = v.split("motherduck_token=")[0] + f"motherduck_token={replace_with}"
         elif isinstance(v, dict):
             remove_passwords_in_dict(v, replace_with)
         elif k.startswith("database"):
@@ -284,14 +287,18 @@ def _join_if_any(sym, args):
 
 
 def remove_password_from_url(url: str, replace_with: str = "***") -> str:
-    parsed = urlparse(url)
-    account = parsed.username or ""
-    if parsed.password:
-        account += ":" + replace_with
-    host = _join_if_any(":", filter(None, [parsed.hostname, parsed.port]))
-    netloc = _join_if_any("@", filter(None, [account, host]))
-    replaced = parsed._replace(netloc=netloc)
-    return replaced.geturl()
+    if "motherduck_token=" in url:
+        replace_token_url = url.split("motherduck_token=")[0] + f"motherduck_token={replace_with}"
+        return replace_token_url
+    else:
+        parsed = urlparse(url)
+        account = parsed.username or ""
+        if parsed.password:
+            account += ":" + replace_with
+        host = _join_if_any(":", filter(None, [parsed.hostname, parsed.port]))
+        netloc = _join_if_any("@", filter(None, [account, host]))
+        replaced = parsed._replace(netloc=netloc)
+        return replaced.geturl()
 
 
 def match_like(pattern: str, strs: Sequence[str]) -> Iterable[str]:

tests/test_joindiff.py

@@ -22,6 +22,7 @@
     db.MySQL,
     db.Snowflake,
     db.BigQuery,
+    db.DuckDB,
     db.Oracle,
     db.Redshift,
     db.Presto,
@@ -32,7 +33,7 @@
 test_each_database = test_each_database_in_list(TEST_DATABASES)
 
 
-@test_each_database_in_list({db.Snowflake, db.BigQuery})
+@test_each_database_in_list({db.Snowflake, db.BigQuery, db.DuckDB})
 class TestCompositeKey(DiffTestCase):
     src_schema = {"id": int, "userid": int, "movieid": int, "rating": float, "timestamp": datetime}
     dst_schema = {"id": int, "userid": int, "movieid": int, "rating": float, "timestamp": datetime}

tests/test_utils.py

@@ -1,6 +1,7 @@
 import unittest
 
 from data_diff.utils import remove_passwords_in_dict, match_regexps, match_like, number_to_human
+from data_diff.__main__ import _remove_passwords_in_dict
 
 
 class TestUtils(unittest.TestCase):
@@ -15,11 +16,48 @@ def test_remove_passwords_in_dict(self):
         remove_passwords_in_dict(d, "$$$$")
         assert d["database_url"] == "mysql://user:$$$$@localhost/db"
 
+        # Test replacing motherduck token in database URL
+        d = {"database_url": "md:datafold_demo?motherduck_token=jaiwefjoaisdk"}
+        remove_passwords_in_dict(d, "$$$$")
+        assert d["database_url"] == "md:datafold_demo?motherduck_token=$$$$"
+
         # Test replacing password in nested dictionary
         d = {"info": {"password": "mypassword"}}
         remove_passwords_in_dict(d, "%%")
         assert d["info"]["password"] == "%%"
 
+        # Test replacing a motherduck token in nested dictionary
+        d = {'database1': {'driver': 'duckdb', 'filepath':'md:datafold_demo?motherduck_token=awieojfaowiejacijobhiwaef'}}
+        remove_passwords_in_dict(d, "%%")
+        assert d["database1"]["filepath"] == "md:datafold_demo?motherduck_token=%%"
+
+    # Test __main__ utility version of this function
+    def test__main__remove_passwords_in_dict(self):
+        # Test replacing password value
+        d = {"password": "mypassword"}
+        _remove_passwords_in_dict(d)
+        assert d["password"] == "**********"
+
+        # Test replacing password in database URL
+        d = {"database_url": "mysql://user:mypassword@localhost/db"}
+        _remove_passwords_in_dict(d)
+        assert d["database_url"] == "mysql://user:***@localhost/db"
+
+        # Test replacing motherduck token in database URL
+        d = {"database_url": "md:datafold_demo?motherduck_token=jaiwefjoaisdk"}
+        _remove_passwords_in_dict(d)
+        assert d["database_url"] == "md:datafold_demo?motherduck_token=***"
+
+        # Test replacing password in nested dictionary
+        d = {"info": {"password": "mypassword"}}
+        _remove_passwords_in_dict(d)
+        assert d["info"]["password"] == "**********"
+
+        # Test replacing a motherduck token in nested dictionary
+        d = {'database1': {'driver': 'duckdb', 'filepath':'md:datafold_demo?motherduck_token=awieojfaowiejacijobhiwaef'}}
+        _remove_passwords_in_dict(d)
+        assert d["database1"]["filepath"] == "md:datafold_demo?motherduck_token=**********"
+
     def test_match_regexps(self):
         def only_results(x):
             return [v for k, v in x]