collate-data-diff 0.11.2__py3-none-any.whl

data_diff/__main__.py

@@ -0,0 +1,618 @@
+import json
+import logging
+import os
+import sys
+import time
+from copy import deepcopy
+from datetime import datetime
+from itertools import islice
+from typing import Dict, Optional, Tuple, Union, List, Set
+
+import click
+import rich
+from rich.logging import RichHandler
+
+from data_diff import Database, DbPath
+from data_diff.config import apply_config_from_file
+from data_diff.databases._connect import connect
+from data_diff.dbt import dbt_diff
+from data_diff.diff_tables import Algorithm, TableDiffer
+from data_diff.hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
+from data_diff.joindiff_tables import TABLE_WRITE_LIMIT, JoinDiffer
+from data_diff.parse_time import parse_time_before, UNITS_STR, ParseError
+from data_diff.queries.api import current_timestamp
+from data_diff.schema import RawColumnInfo, create_schema
+from data_diff.table_segment import TableSegment
+from data_diff.tracking import disable_tracking, set_entrypoint_name
+from data_diff.utils import eval_name_template, remove_password_from_url, safezip, match_like, LogStatusHandler
+from data_diff.version import __version__
+
+COLOR_SCHEME = {
+    "+": "green",
+    "-": "red",
+}
+
+set_entrypoint_name(os.getenv("DATAFOLD_TRIGGERED_BY", "CLI"))
+
+
+def _get_log_handlers(is_dbt: Optional[bool] = False) -> Dict[str, logging.Handler]:
+    handlers = {}
+    date_format = "%H:%M:%S"
+    log_format_rich = "%(message)s"
+
+    # limits to 100 characters arbitrarily
+    log_format_status = "%(message).100s"
+    rich_handler = RichHandler(rich_tracebacks=True)
+    rich_handler.setFormatter(logging.Formatter(log_format_rich, datefmt=date_format))
+    rich_handler.setLevel(logging.WARN)
+    handlers["rich_handler"] = rich_handler
+
+    # only use log_status_handler in an interactive terminal session
+    if rich_handler.console.is_interactive and is_dbt:
+        log_status_handler = LogStatusHandler()
+        log_status_handler.setFormatter(logging.Formatter(log_format_status, datefmt=date_format))
+        log_status_handler.setLevel(logging.DEBUG)
+        handlers["log_status_handler"] = log_status_handler
+
+    return handlers
+
+
+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"):
+            d[k] = remove_password_from_url(v)
+
+
+def _get_schema(pair: Tuple[Database, DbPath]) -> Dict[str, RawColumnInfo]:
+    db, table_path = pair
+    return db.query_table_schema(table_path)
+
+
+def diff_schemas(table1, table2, schema1, schema2, columns) -> None:
+    logging.info("Diffing schemas...")
+    attrs = "name", "type", "datetime_precision", "numeric_precision", "numeric_scale"
+    for c in columns:
+        if c is None:  # Skip for convenience
+            continue
+        diffs = []
+
+        if c not in schema1:
+            cols = ", ".join(schema1)
+            raise ValueError(f"Column '{c}' not found in table 1, named '{table1}'. Columns: {cols}")
+        if c not in schema2:
+            cols = ", ".join(schema1)
+            raise ValueError(f"Column '{c}' not found in table 2, named '{table2}'. Columns: {cols}")
+
+        col1 = schema1[c]
+        col2 = schema2[c]
+
+        for attr, v1, v2 in safezip(attrs, col1, col2):
+            if v1 != v2:
+                diffs.append(f"{attr}:({v1} != {v2})")
+        if diffs:
+            logging.warning(f"Schema mismatch in column '{c}': {', '.join(diffs)}")
+
+
+class MyHelpFormatter(click.HelpFormatter):
+    def __init__(self, **kwargs) -> None:
+        super().__init__(self, **kwargs)
+        self.indent_increment = 6
+
+    def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -> None:
+        self.write(f"data-diff v{__version__} - efficiently diff rows across database tables.\n\n")
+        self.write("Usage:\n")
+        self.write(f"  * In-db diff:    {prog} <database_a> <table_a> <table_b> [OPTIONS]\n")
+        self.write(f"  * Cross-db diff: {prog} <database_a> <table_a> <database_b> <table_b> [OPTIONS]\n")
+        self.write(f"  * Using config:  {prog} --conf PATH [--run NAME] [OPTIONS]\n")
+
+
+click.Context.formatter_class = MyHelpFormatter
+
+
+@click.command(no_args_is_help=True)
+@click.argument("database1", required=False)
+@click.argument("table1", required=False)
+@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"
+)
+@click.option("-t", "--update-column", default=None, help="Name of updated_at/last_updated column", metavar="NAME")
+@click.option(
+    "-c",
+    "--columns",
+    default=[],
+    multiple=True,
+    help="Names of extra columns to compare."
+    "Can be used more than once in the same command. "
+    "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(
+    "--bisection-factor",
+    default=None,
+    help=f"Segments per iteration. Default={DEFAULT_BISECTION_FACTOR}.",
+    metavar="NUM",
+)
+@click.option(
+    "--bisection-threshold",
+    default=None,
+    help=f"Minimal bisection threshold. Below it, data-diff will download the data and compare it locally. Default={DEFAULT_BISECTION_THRESHOLD}.",
+    metavar="NUM",
+)
+@click.option(
+    "-m",
+    "--materialize-to-table",
+    default=None,
+    metavar="TABLE_NAME",
+    help="(joindiff only) Materialize the diff results into a new table in the database. If a table exists by that name, it will be replaced.",
+)
+@click.option(
+    "--min-age",
+    default=None,
+    help="Considers only rows older than specified. Useful for specifying replication lag."
+    "Example: --min-age=5min ignores rows from the last 5 minutes. "
+    f"\nValid units: {UNITS_STR}",
+    metavar="AGE",
+)
+@click.option(
+    "--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("-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(
+    "--case-sensitive",
+    is_flag=True,
+    help="Column names are treated as case-sensitive. Otherwise, data-diff corrects their case according to schema.",
+)
+@click.option(
+    "--assume-unique-key",
+    is_flag=True,
+    help="Skip validating the uniqueness of the key column during joindiff, which is costly in non-cloud dbs.",
+)
+@click.option(
+    "--sample-exclusive-rows",
+    is_flag=True,
+    help="Sample several rows that only appear in one of the tables, but not the other. (joindiff only)",
+)
+@click.option(
+    "--materialize-all-rows",
+    is_flag=True,
+    help="Materialize every row, even if they are the same, instead of just the differing rows. (joindiff only)",
+)
+@click.option(
+    "--table-write-limit",
+    default=TABLE_WRITE_LIMIT,
+    help=f"Maximum number of rows to write when creating materialized or sample tables, per thread. Default={TABLE_WRITE_LIMIT}",
+    metavar="COUNT",
+)
+@click.option(
+    "-j",
+    "--threads",
+    default=None,
+    help="Number of worker threads to use per database. Default=1. "
+    "A higher number will increase performance, but take more capacity from your database. "
+    "'serial' guarantees a single-threaded execution of the algorithm (useful for debugging).",
+    metavar="COUNT",
+)
+@click.option(
+    "-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(
+    "--conf",
+    default=None,
+    help="Path to a configuration.toml file, to provide a default configuration, and a list of possible runs.",
+    metavar="PATH",
+)
+@click.option(
+    "--run",
+    default=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 a diff using your local dbt project. Expects to be run from a dbt project folder by default.",
+)
+@click.option(
+    "--cloud",
+    is_flag=True,
+    help="Add this flag along with --dbt to run a diff using your local dbt project on Datafold cloud. Expects an api key on env var DATAFOLD_API_KEY.",
+)
+@click.option(
+    "--dbt-profiles-dir",
+    envvar="DBT_PROFILES_DIR",
+    default=None,
+    metavar="PATH",
+    help="Which directory to look in for the profiles.yml file. If not set, we follow the default profiles.yml location for the dbt version being used. Can also be set via the DBT_PROFILES_DIR environment variable.",
+)
+@click.option(
+    "--dbt-project-dir",
+    default=None,
+    metavar="PATH",
+    help="Which directory to look in for the dbt_project.yml file. Default is the current working directory and its parents.",
+)
+@click.option(
+    "--select",
+    "-s",
+    default=None,
+    metavar="SELECTION or MODEL_NAME",
+    help="--select dbt resources to compare using dbt selection syntax in dbt versions >= 1.5.\nIn versions < 1.5, it will naively search for a model with MODEL_NAME as the name.",
+)
+@click.option(
+    "--state",
+    "-s",
+    default=None,
+    metavar="PATH",
+    help="Specify manifest to utilize for 'prod' comparison paths instead of using configuration.",
+)
+@click.option(
+    "-pd",
+    "--prod-database",
+    "prod_database",
+    default=None,
+    help="Override the dbt production database configuration within dbt_project.yml",
+)
+@click.option(
+    "-ps",
+    "--prod-schema",
+    "prod_schema",
+    default=None,
+    help="Override the dbt production schema configuration within dbt_project.yml",
+)
+def main(conf, run, **kw) -> None:
+    log_handlers = _get_log_handlers(kw["dbt"])
+    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"]:
+        log_handlers["rich_handler"].setLevel(logging.DEBUG)
+        logging.basicConfig(level=logging.DEBUG, handlers=list(log_handlers.values()))
+        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"):
+        log_handlers["rich_handler"].setLevel(logging.INFO)
+        logging.basicConfig(level=logging.DEBUG, handlers=list(log_handlers.values()))
+    else:
+        log_handlers["rich_handler"].setLevel(logging.WARNING)
+        logging.basicConfig(level=logging.DEBUG, handlers=list(log_handlers.values()))
+
+    try:
+        state = kw.pop("state", None)
+        if state:
+            state = os.path.expanduser(state)
+        profiles_dir_override = kw.pop("dbt_profiles_dir", None)
+        if profiles_dir_override:
+            profiles_dir_override = os.path.expanduser(profiles_dir_override)
+        project_dir_override = kw.pop("dbt_project_dir", None)
+        if project_dir_override:
+            project_dir_override = os.path.expanduser(project_dir_override)
+        if kw["dbt"]:
+            dbt_diff(
+                log_status_handler=log_handlers.get("log_status_handler"),
+                profiles_dir_override=profiles_dir_override,
+                project_dir_override=project_dir_override,
+                is_cloud=kw["cloud"],
+                dbt_selection=kw["select"],
+                json_output=kw["json_output"],
+                state=state,
+                where_flag=kw["where"],
+                stats_flag=kw["stats"],
+                columns_flag=kw["columns"],
+                production_database_flag=kw["prod_database"],
+                production_schema_flag=kw["prod_schema"],
+            )
+        else:
+            _data_diff(dbt_project_dir=project_dir_override, dbt_profiles_dir=profiles_dir_override, state=state, **kw)
+    except Exception as e:
+        logging.error(e)
+        raise
+
+
+def _get_dbs(
+    threads: int, database1: str, threads1: int, database2: str, threads2: int, interactive: bool
+) -> Tuple[Database, Database]:
+    db1 = connect(database1, threads1 or threads)
+    if database1 == database2:
+        db2 = db1
+    else:
+        db2 = connect(database2, threads2 or threads)
+
+    if interactive:
+        db1.enable_interactive()
+        db2.enable_interactive()
+
+    return db1, db2
+
+
+def _set_age(options: dict, min_age: Optional[str], max_age: Optional[str], db: Database) -> None:
+    if min_age or max_age:
+        now: datetime = db.query(current_timestamp(), datetime).replace(tzinfo=None)
+        try:
+            if max_age:
+                options["min_update"] = parse_time_before(now, max_age)
+            if min_age:
+                options["max_update"] = parse_time_before(now, min_age)
+        except ParseError as e:
+            logging.error(f"Error while parsing age expression: {e}")
+
+
+def _get_table_differ(
+    algorithm: str,
+    db1: Database,
+    db2: Database,
+    threaded: bool,
+    threads: int,
+    assume_unique_key: bool,
+    sample_exclusive_rows: bool,
+    materialize_all_rows: bool,
+    table_write_limit: int,
+    materialize_to_table: Optional[str],
+    bisection_factor: Optional[int],
+    bisection_threshold: Optional[int],
+) -> TableDiffer:
+    algorithm = Algorithm(algorithm)
+    if algorithm == Algorithm.AUTO:
+        algorithm = Algorithm.JOINDIFF if db1 == db2 else Algorithm.HASHDIFF
+
+    logging.info(f"Using algorithm '{algorithm.name.lower()}'.")
+
+    if algorithm == Algorithm.JOINDIFF:
+        return JoinDiffer(
+            threaded=threaded,
+            max_threadpool_size=threads and threads * 2,
+            validate_unique_key=not assume_unique_key,
+            sample_exclusive_rows=sample_exclusive_rows,
+            materialize_all_rows=materialize_all_rows,
+            table_write_limit=table_write_limit,
+            materialize_to_table=(
+                materialize_to_table and db1.dialect.parse_table_name(eval_name_template(materialize_to_table))
+            ),
+        )
+
+    assert algorithm == Algorithm.HASHDIFF
+    return HashDiffer(
+        bisection_factor=DEFAULT_BISECTION_FACTOR if bisection_factor is None else bisection_factor,
+        bisection_threshold=DEFAULT_BISECTION_THRESHOLD if bisection_threshold is None else bisection_threshold,
+        threaded=threaded,
+        max_threadpool_size=threads and threads * 2,
+    )
+
+
+def _print_result(stats, json_output, diff_iter) -> None:
+    if stats:
+        if json_output:
+            rich.print(json.dumps(diff_iter.get_stats_dict()))
+        else:
+            rich.print(diff_iter.get_stats_string())
+
+    else:
+        for op, values in diff_iter:
+            color = COLOR_SCHEME.get(op, "grey62")
+
+            if json_output:
+                jsonl = json.dumps([op, list(values)])
+                rich.print(f"[{color}]{jsonl}[/{color}]")
+            else:
+                text = f"{op} {', '.join(map(str, values))}"
+                rich.print(f"[{color}]{text}[/{color}]")
+
+            sys.stdout.flush()
+
+
+def _get_expanded_columns(
+    columns: List[str],
+    case_sensitive: bool,
+    mutual: Set[str],
+    db1: Database,
+    schema1: dict,
+    table1: str,
+    db2: Database,
+    schema2: dict,
+    table2: str,
+) -> Set[str]:
+    expanded_columns: Set[str] = set()
+    for c in columns:
+        cc = c if case_sensitive else c.lower()
+        match = set(match_like(cc, mutual))
+        if not match:
+            m1 = None if any(match_like(cc, schema1.keys())) else f"{db1}/{table1}"
+            m2 = None if any(match_like(cc, schema2.keys())) else f"{db2}/{table2}"
+            not_matched = ", ".join(m for m in [m1, m2] if m)
+            raise ValueError(f"Column '{c}' not found in: {not_matched}")
+
+        expanded_columns |= match
+    return expanded_columns
+
+
+def _get_threads(threads: Union[int, str, None], threads1: Optional[int], threads2: Optional[int]) -> Tuple[bool, int]:
+    threaded = True
+    if threads is None:
+        threads = 1
+    elif isinstance(threads, str) and threads.lower() == "serial":
+        assert not (threads1 or threads2)
+        threaded = False
+        threads = 1
+    else:
+        try:
+            threads = int(threads)
+        except ValueError:
+            logging.error("Error: threads must be a number, or 'serial'.")
+            raise
+
+        if threads < 1:
+            logging.error("Error: threads must be >= 1")
+            raise ValueError("Error: threads must be >= 1")
+
+    return threaded, threads
+
+
+def _data_diff(
+    database1,
+    table1,
+    database2,
+    table2,
+    key_columns,
+    update_column,
+    columns,
+    limit,
+    algorithm,
+    bisection_factor,
+    bisection_threshold,
+    min_age,
+    max_age,
+    stats,
+    debug,
+    verbose,
+    version,
+    interactive,
+    no_tracking,
+    threads,
+    case_sensitive,
+    json_output,
+    where,
+    assume_unique_key,
+    sample_exclusive_rows,
+    materialize_all_rows,
+    table_write_limit,
+    materialize_to_table,
+    dbt,
+    cloud,
+    dbt_profiles_dir,
+    dbt_project_dir,
+    prod_database,
+    prod_schema,
+    select,
+    state,
+    threads1=None,
+    threads2=None,
+    __conf__=None,
+) -> None:
+    if limit and stats:
+        logging.error("Cannot specify a limit when using the -s/--stats switch")
+        return
+
+    key_columns = key_columns or ("id",)
+    threaded, threads = _get_threads(threads, threads1, threads2)
+    start = time.monotonic()
+
+    if database1 is None or database2 is None:
+        logging.error(
+            f"Error: Databases not specified. Got {database1} and {database2}. Use --help for more information."
+        )
+        return
+
+    db1: Database
+    db2: Database
+    db1, db2 = _get_dbs(threads, database1, threads1, database2, threads2, interactive)
+    with db1, db2:
+        options = {
+            "case_sensitive": case_sensitive,
+            "where": where,
+        }
+
+        _set_age(options, min_age, max_age, db1)
+        dbs: Tuple[Database, Database] = db1, db2
+
+        differ = _get_table_differ(
+            algorithm,
+            db1,
+            db2,
+            threaded,
+            threads,
+            assume_unique_key,
+            sample_exclusive_rows,
+            materialize_all_rows,
+            table_write_limit,
+            materialize_to_table,
+            bisection_factor,
+            bisection_threshold,
+        )
+
+        table_names = table1, table2
+        table_paths = [db.dialect.parse_table_name(t) for db, t in safezip(dbs, table_names)]
+
+        schemas = list(differ._thread_map(_get_schema, safezip(dbs, table_paths)))
+        schema1, schema2 = schemas = [
+            create_schema(db.name, table_path, schema, case_sensitive)
+            for db, table_path, schema in safezip(dbs, table_paths, schemas)
+        ]
+
+        mutual = schema1.keys() & schema2.keys()  # Case-aware, according to case_sensitive
+        logging.debug(f"Available mutual columns: {mutual}")
+
+        expanded_columns = _get_expanded_columns(
+            columns, case_sensitive, mutual, db1, schema1, table1, db2, schema2, table2
+        )
+        columns = tuple(expanded_columns - {*key_columns, update_column})
+
+        if db1 == db2:
+            diff_schemas(
+                table_names[0],
+                table_names[1],
+                schema1,
+                schema2,
+                (
+                    *key_columns,
+                    update_column,
+                    *columns,
+                ),
+            )
+
+        logging.info(f"Diffing using columns: key={key_columns} update={update_column} extra={columns}.")
+
+        segments = [
+            TableSegment(db, table_path, key_columns, update_column, columns, **options)._with_raw_schema(raw_schema)
+            for db, table_path, raw_schema in safezip(dbs, table_paths, schemas)
+        ]
+
+        diff_iter = differ.diff_tables(*segments)
+
+        if limit:
+            assert not stats
+            diff_iter = islice(diff_iter, int(limit))
+
+        _print_result(stats, json_output, diff_iter)
+
+    end = time.monotonic()
+    logging.info(f"Duration: {end-start:.2f} seconds.")
+
+
+if __name__ == "__main__":
+    main()

data_diff/abcs/compiler.py

@@ -0,0 +1,13 @@
+from abc import ABC
+
+import attrs
+
+
+@attrs.define(frozen=False)
+class AbstractCompiler(ABC):
+    pass
+
+
+@attrs.define(frozen=False, eq=False)
+class Compilable(ABC):
+    pass