squash add dbt integration

Author: dlawin

data_diff/__main__.py

@@ -13,9 +13,14 @@
 from sqeleton.schema import create_schema
 from sqeleton.queries.api import current_timestamp
 
+from .dbt import DbtDiffer
 from .utils import eval_name_template, remove_password_from_url, safezip, match_like
 from .diff_tables import Algorithm
-from .hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
+from .hashdiff_tables import (
+    HashDiffer,
+    DEFAULT_BISECTION_THRESHOLD,
+    DEFAULT_BISECTION_FACTOR,
+)
 from .joindiff_tables import TABLE_WRITE_LIMIT, JoinDiffer
 from .table_segment import TableSegment
 from .databases import connect
@@ -98,9 +103,20 @@ def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -
 @click.argument("database2", required=False)
 @click.argument("table2", required=False)
 @click.option(
-    "-k", "--key-columns", default=[], multiple=True, help="Names of primary key columns. Default='id'.", metavar="NAME"
+    "-k",
+    "--key-columns",
+    default=[],
+    multiple=True,
+    help="Names of primary key columns. Default='id'.",
+    metavar="NAME",
+)
+@click.option(
+    "-t",
+    "--update-column",
+    default=None,
+    help="Name of updated_at/last_updated column",
+    metavar="NAME",
 )
-@click.option("-t", "--update-column", default=None, help="Name of updated_at/last_updated column", metavar="NAME")
 @click.option(
     "-c",
     "--columns",
@@ -111,7 +127,13 @@ def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -
     "Accepts a name or a pattern like in SQL. Example: -c col% -c another_col",
     metavar="NAME",
 )
-@click.option("-l", "--limit", default=None, help="Maximum number of differences to find", metavar="NUM")
+@click.option(
+    "-l",
+    "--limit",
+    default=None,
+    help="Maximum number of differences to find",
+    metavar="NUM",
+)
 @click.option(
     "--bisection-factor",
     default=None,
@@ -140,15 +162,27 @@ def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -
     metavar="AGE",
 )
 @click.option(
-    "--max-age", default=None, help="Considers only rows younger than specified. See --min-age.", metavar="AGE"
+    "--max-age",
+    default=None,
+    help="Considers only rows younger than specified. See --min-age.",
+    metavar="AGE",
 )
 @click.option("-s", "--stats", is_flag=True, help="Print stats instead of a detailed diff")
 @click.option("-d", "--debug", is_flag=True, help="Print debug info")
-@click.option("--json", "json_output", is_flag=True, help="Print JSONL output for machine readability")
+@click.option(
+    "--json",
+    "json_output",
+    is_flag=True,
+    help="Print JSONL output for machine readability",
+)
 @click.option("-v", "--verbose", is_flag=True, help="Print extra info")
 @click.option("--version", is_flag=True, help="Print version info and exit")
 @click.option("-i", "--interactive", is_flag=True, help="Confirm queries, implies --debug")
-@click.option("--no-tracking", is_flag=True, help="data-diff sends home anonymous usage data. Use this to disable it.")
+@click.option(
+    "--no-tracking",
+    is_flag=True,
+    help="data-diff sends home anonymous usage data. Use this to disable it.",
+)
 @click.option(
     "--case-sensitive",
     is_flag=True,
@@ -185,9 +219,18 @@ def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -
     metavar="COUNT",
 )
 @click.option(
-    "-w", "--where", default=None, help="An additional 'where' expression to restrict the search space. Beware of SQL Injection!", metavar="EXPR"
+    "-w",
+    "--where",
+    default=None,
+    help="An additional 'where' expression to restrict the search space. Beware of SQL Injection!",
+    metavar="EXPR",
+)
+@click.option(
+    "-a",
+    "--algorithm",
+    default=Algorithm.AUTO.value,
+    type=click.Choice([i.value for i in Algorithm]),
 )
-@click.option("-a", "--algorithm", default=Algorithm.AUTO.value, type=click.Choice([i.value for i in Algorithm]))
 @click.option(
     "--conf",
     default=None,
@@ -200,24 +243,74 @@ def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -
     help="Name of run-configuration to run. If used, CLI arguments for database and table must be omitted.",
     metavar="NAME",
 )
+@click.option(
+    "--dbt",
+    is_flag=True,
+    help="Run dbt diff.",
+)
+@click.option(
+    "--dbt-cloud",
+    is_flag=True,
+    help="Run dbt diff on Datafold cloud. Expects an api key on env var DATAFOLD_API_KEY.",
+)
+@click.option(
+    "--dbt-profiles-dir",
+    default=None,
+    metavar="PATH",
+    help="Override the default dbt profile location (~/.dbt).",
+)
+@click.option(
+    "--dbt-project-dir",
+    default=None,
+    metavar="PATH",
+    help="Override the dbt project directory. Otherwise assumed to be the current directory.",
+)
 def main(conf, run, **kw):
     if kw["table2"] is None and kw["database2"]:
         # Use the "database table table" form
         kw["table2"] = kw["database2"]
         kw["database2"] = kw["database1"]
 
+    if kw["version"]:
+        print(f"v{__version__}")
+        return
+
     if conf:
         kw = apply_config_from_file(conf, run, kw)
 
+    if kw["no_tracking"]:
+        disable_tracking()
+
+    if kw.get("interactive"):
+        kw["debug"] = True
+
+    if kw["debug"]:
+        logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT, datefmt=DATE_FORMAT)
+        if kw.get("__conf__"):
+            kw["__conf__"] = deepcopy(kw["__conf__"])
+            _remove_passwords_in_dict(kw["__conf__"])
+            logging.debug(f"Applied run configuration: {kw['__conf__']}")
+    elif kw.get("verbose"):
+        logging.basicConfig(level=logging.INFO, format=LOG_FORMAT, datefmt=DATE_FORMAT)
+    else:
+        logging.basicConfig(level=logging.WARNING, format=LOG_FORMAT, datefmt=DATE_FORMAT)
+
     try:
-        return _main(**kw)
+        if kw["dbt"] or kw["dbt_cloud"]:
+            DbtDiffer.diff(
+                profiles_dir_override=kw["dbt_profiles_dir"],
+                project_dir_override=kw["dbt_project_dir"],
+                is_cloud=kw["dbt_cloud"],
+            )
+        else:
+            return _data_diff(**kw)
     except Exception as e:
         logging.error(e)
         if kw["debug"]:
             raise
 
 
-def _main(
+def _data_diff(
     database1,
     table1,
     database2,
@@ -246,31 +339,14 @@ def _main(
     materialize_all_rows,
     table_write_limit,
     materialize_to_table,
+    dbt,
+    dbt_cloud,
+    dbt_profiles_dir,
+    dbt_project_dir,
     threads1=None,
     threads2=None,
     __conf__=None,
 ):
-    if version:
-        print(f"v{__version__}")
-        return
-
-    if no_tracking:
-        disable_tracking()
-
-    if interactive:
-        debug = True
-
-    if debug:
-        logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT, datefmt=DATE_FORMAT)
-        if __conf__:
-            __conf__ = deepcopy(__conf__)
-            _remove_passwords_in_dict(__conf__)
-            logging.debug(f"Applied run configuration: {__conf__}")
-    elif verbose:
-        logging.basicConfig(level=logging.INFO, format=LOG_FORMAT, datefmt=DATE_FORMAT)
-    else:
-        logging.basicConfig(level=logging.WARNING, format=LOG_FORMAT, datefmt=DATE_FORMAT)
-
     if limit and stats:
         logging.error("Cannot specify a limit when using the -s/--stats switch")
         return