dcs-sdk 1.6.5__py3-none-any.whl

data_diff/__main__.py

@@ -0,0 +1,517 @@
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+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, List, Optional, Set, Tuple, Union
+
+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.diff_tables import Algorithm, TableDiffer
+from data_diff.hashdiff_tables import (
+    DEFAULT_BISECTION_FACTOR,
+    DEFAULT_BISECTION_THRESHOLD,
+    HashDiffer,
+)
+from data_diff.joindiff_tables import TABLE_WRITE_LIMIT, JoinDiffer
+from data_diff.parse_time import UNITS_STR, ParseError, parse_time_before
+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.utils import (
+    eval_name_template,
+    match_like,
+    remove_password_from_url,
+    safezip,
+    set_entrypoint_name,
+)
+from data_diff.version import __version__
+
+COLOR_SCHEME = {
+    "+": "green",
+    "-": "red",
+}
+
+set_entrypoint_name(os.getenv("DATADIFF_TRIGGERED_BY", "CLI"))
+
+
+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"dcs-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, dcs-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(
+    "--case-sensitive",
+    is_flag=True,
+    help="Column names are treated as case-sensitive. Otherwise, dcs-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",
+)
+def main(conf, run, **kw) -> None:
+    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
+
+    try:
+        _data_diff(**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,
+    threads,
+    case_sensitive,
+    json_output,
+    where,
+    assume_unique_key,
+    sample_exclusive_rows,
+    materialize_all_rows,
+    table_write_limit,
+    materialize_to_table,
+    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/__init__.py

@@ -0,0 +1,13 @@
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.

data_diff/abcs/compiler.py

@@ -0,0 +1,27 @@
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from abc import ABC
+
+import attrs
+
+
+@attrs.define(frozen=False)
+class AbstractCompiler(ABC):
+    pass
+
+
+@attrs.define(frozen=False, eq=False)
+class Compilable(ABC):
+    pass