From 5dec5e158c40e1ab5c1b561aff767d59f8a403cb Mon Sep 17 00:00:00 2001 From: Matteo Renoldi Date: Sat, 21 Feb 2026 09:48:26 +0100 Subject: [PATCH] feat: add a cli count command --- README.md | 15 + docs/index.md | 15 + src/quack_diff/cli/commands/__init__.py | 3 +- src/quack_diff/cli/commands/count/__init__.py | 326 ++++++++++++++++++ src/quack_diff/cli/formatters.py | 48 ++- src/quack_diff/cli/main.py | 3 +- src/quack_diff/cli/output.py | 57 ++- src/quack_diff/core/differ.py | 84 +++++ src/quack_diff/core/query_builder.py | 38 ++ tests/cli/test_main.py | 78 +++++ tests/conftest.py | 46 +++ tests/core/test_differ.py | 123 +++++++ tests/core/test_query_builder.py | 23 ++ 13 files changed, 855 insertions(+), 4 deletions(-) create mode 100644 src/quack_diff/cli/commands/count/__init__.py diff --git a/README.md b/README.md index 89b7025..df73de5 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,21 @@ quack-diff compare \ quack-diff schema --source db1.users --target db2.users ``` +### Count check (bronze/silver/gold) + +Validate that pipeline layers have the same number of rows or distinct keys (no row loss): + +```bash +# Same row count across layers +quack-diff count -t bronze.orders -t silver.orders -t gold.orders + +# Same distinct ID count +quack-diff count -t bronze.orders -t silver.orders -t gold.orders --key order_id + +# JSON for CI/CD +quack-diff count -t bronze.orders -t silver.orders -t gold.orders --key order_id --json +``` + ## Configuration quack-diff supports configuration via environment variables: diff --git a/docs/index.md b/docs/index.md index 077cd51..ae4e815 100644 --- a/docs/index.md +++ b/docs/index.md @@ -42,6 +42,21 @@ quack-diff compare \ quack-diff schema --source db1.users --target db2.users ``` +### Count check (bronze/silver/gold) + +Validate that pipeline layers have the same number of rows or distinct keys (no row loss): + +```bash +# Same row count across layers +quack-diff count -t bronze.orders -t silver.orders -t gold.orders + +# Same distinct ID count +quack-diff count -t bronze.orders -t silver.orders -t gold.orders --key order_id + +# JSON for CI/CD +quack-diff count -t bronze.orders -t silver.orders -t gold.orders --key order_id --json +``` + ## How It Works quack-diff leverages DuckDB's extension system to connect to external databases: diff --git a/src/quack_diff/cli/commands/__init__.py b/src/quack_diff/cli/commands/__init__.py index 43f7c57..023ce48 100644 --- a/src/quack_diff/cli/commands/__init__.py +++ b/src/quack_diff/cli/commands/__init__.py @@ -2,6 +2,7 @@ from quack_diff.cli.commands.attach import attach from quack_diff.cli.commands.compare import compare +from quack_diff.cli.commands.count import count from quack_diff.cli.commands.schema import schema -__all__ = ["attach", "compare", "schema"] +__all__ = ["attach", "compare", "count", "schema"] diff --git a/src/quack_diff/cli/commands/count/__init__.py b/src/quack_diff/cli/commands/count/__init__.py new file mode 100644 index 0000000..e4652eb --- /dev/null +++ b/src/quack_diff/cli/commands/count/__init__.py @@ -0,0 +1,326 @@ +"""Count command for validating row counts across multiple tables.""" + +from __future__ import annotations + +import logging +import time +from pathlib import Path +from typing import TYPE_CHECKING, Annotated + +import typer + +from quack_diff.cli.console import ( + console, + print_error, + print_success, + set_json_output_mode, + status_context, +) +from quack_diff.cli.errors import get_error_info +from quack_diff.cli.formatters import print_count_result +from quack_diff.cli.output import format_count_result_json, format_error_json, print_json +from quack_diff.config import get_settings +from quack_diff.core.connector import DuckDBConnector +from quack_diff.core.differ import DataDiffer +from quack_diff.core.sql_utils import ( + AttachError, + DatabaseError, + QueryExecutionError, + SQLInjectionError, + TableNotFoundError, +) + +if TYPE_CHECKING: + from quack_diff.config import Settings + +logger = logging.getLogger(__name__) + + +def _parse_table_reference(table: str, known_aliases: set[str] | None = None) -> tuple[str | None, str]: + """Parse a table reference to extract alias and table name. + + If known_aliases is provided (e.g. from settings.databases), any first + segment in that set is treated as an alias. + """ + parts = table.split(".", 1) + if len(parts) == 2 and parts[0].lower() in ("sf", "snowflake"): + return parts[0].lower(), parts[1] + if len(parts) >= 2: + first_part = parts[0].lower() + if known_aliases and first_part in known_aliases: + return first_part, ".".join(parts[1:]) + if len(first_part) <= 4 and first_part.isalpha(): + return first_part, ".".join(parts[1:]) + return None, table + + +def _is_snowflake_table(table: str, settings: Settings) -> bool: + """Check if a table reference points to a Snowflake table.""" + known_aliases = set(settings.databases.keys()) if settings.databases else set() + alias, _ = _parse_table_reference(table, known_aliases=known_aliases) + if alias in ("sf", "snowflake"): + return True + if alias and alias in settings.databases: + db_config = settings.databases[alias] + return db_config.get("type", "snowflake").lower() == "snowflake" + return False + + +def _auto_attach_databases( + connector: DuckDBConnector, + settings: Settings, + tables: list[str], + verbose: bool = False, +) -> None: + """Auto-attach DuckDB databases for any aliases in table references.""" + known_aliases = set(settings.databases.keys()) if settings.databases else set() + aliases_to_attach: set[str] = set() + for table in tables: + alias, _ = _parse_table_reference(table, known_aliases=known_aliases) + if alias: + aliases_to_attach.add(alias) + + for alias in aliases_to_attach: + if alias in connector.attached_databases: + continue + if alias in settings.databases: + db_config = settings.databases[alias] + db_type = db_config.get("type", "duckdb").lower() + if db_type == "duckdb": + path = db_config.get("path") + if path: + if verbose: + from quack_diff.cli.console import print_info + + print_info(f"Attaching DuckDB database: {path} as '{alias}'") + connector.attach_duckdb(alias, str(path)) + + +def _resolve_tables_for_count( + connector: DuckDBConnector, + settings: Settings, + tables: list[str], + verbose: bool = False, +) -> tuple[list[str], dict[str, str]]: + """Resolve table list: pull Snowflake tables locally, return (resolved_names, display_name_map).""" + resolved: list[str] = [] + display_name_map: dict[str, str] = {} + + known_aliases = set(settings.databases.keys()) if settings.databases else set() + for i, table in enumerate(tables): + alias, table_name = _parse_table_reference(table, known_aliases=known_aliases) + if alias and _is_snowflake_table(table, settings): + local_name = f"__count_pulled_{i}" + if verbose: + from quack_diff.cli.console import print_info + + print_info(f"Pulling Snowflake table: {table_name}") + + config = None + database = None + if alias in settings.databases: + db_config = settings.databases[alias] + connection_name = db_config.get("connection_name") + database = db_config.get("database") + if connection_name: + from quack_diff.config import SnowflakeConfig + + config = SnowflakeConfig(connection_name=connection_name) + if config is None: + config = settings.snowflake + + connector.pull_snowflake_table( + table_name=table_name, + local_name=local_name, + config=config, + database=database, + ) + resolved.append(local_name) + display_name_map[local_name] = table + else: + resolved.append(table) + display_name_map[table] = table + + return resolved, display_name_map + + +def count( + tables: Annotated[ + list[str], + typer.Option( + "--tables", + "-t", + help="Table(s) to compare counts (repeat or comma-separated)", + ), + ], + key: Annotated[ + str | None, + typer.Option( + "--key", + "-k", + help="Column for COUNT(DISTINCT ...); omit for COUNT(*)", + ), + ] = None, + config_file: Annotated[ + Path | None, + typer.Option( + "--config", + help="Path to configuration file (YAML)", + ), + ] = None, + verbose: Annotated[ + bool, + typer.Option( + "--verbose", + help="Show detailed output", + ), + ] = False, + json_output: Annotated[ + bool, + typer.Option( + "--json", + help="Output results as JSON for CI/CD integration", + ), + ] = False, +) -> None: + """Check that multiple tables have the same row or distinct-key count. + + Use this to validate bronze/silver/gold (or any pipeline) layers + have the same number of records without running a full diff. + + Examples: + + # Same row count across layers + quack-diff count -t bronze.orders -t silver.orders -t gold.orders + + # Same distinct ID count + quack-diff count -t bronze.orders -t silver.orders -t gold.orders --key order_id + + # Comma-separated tables + quack-diff count --tables bronze.orders,silver.orders,gold.orders --key id + + # JSON for CI/CD + quack-diff count -t bronze.orders -t silver.orders -t gold.orders --key id --json + """ + if json_output: + set_json_output_mode(True) + + start_time = time.time() + + # Flatten: support both -t a -t b and -t "a,b" + flat_tables: list[str] = [] + for t in tables: + flat_tables.extend(s.strip() for s in t.split(",") if s.strip()) + + if len(flat_tables) < 2: + if json_output: + print_json( + format_error_json( + error_type="ValueError", + message="At least two tables are required for count check", + exit_code=2, + ) + ) + else: + print_error("At least two tables are required for count check") + raise typer.Exit(2) + + try: + settings = get_settings(config_file=config_file) + use_snowflake = any(_is_snowflake_table(t, settings) for t in flat_tables) + + with DuckDBConnector(settings=settings) as connector: + if use_snowflake: + with status_context("Pulling data from Snowflake..."): + resolved_tables, display_name_map = _resolve_tables_for_count( + connector=connector, + settings=settings, + tables=flat_tables, + verbose=verbose, + ) + else: + _auto_attach_databases(connector, settings, flat_tables, verbose) + resolved_tables = flat_tables + display_name_map = {t: t for t in flat_tables} + + differ = DataDiffer( + connector=connector, + null_sentinel=settings.defaults.null_sentinel, + column_delimiter=settings.defaults.column_delimiter, + ) + + with status_context("Counting..."): + result = differ.count_check( + tables=resolved_tables, + key_column=key, + ) + + duration = time.time() - start_time + + if json_output: + json_data = format_count_result_json( + result, + display_name_map=display_name_map, + duration_seconds=duration, + ) + print_json(json_data) + else: + print_count_result(result, display_name_map=display_name_map) + + if result.is_match: + print_success("All table counts match!") + raise typer.Exit(0) + else: + print_error("Table counts do not match") + raise typer.Exit(1) + + except typer.Exit: + raise + except TableNotFoundError as e: + _handle_error(e, "Table not found", verbose, json_output, start_time) + except QueryExecutionError as e: + _handle_error(e, "Query execution error", verbose, json_output, start_time) + except SQLInjectionError as e: + _handle_error(e, "Invalid input", verbose, json_output, start_time) + except AttachError as e: + _handle_error(e, "Database attach error", verbose, json_output, start_time) + except DatabaseError as e: + _handle_error(e, "Database error", verbose, json_output, start_time) + except ValueError as e: + _handle_error(e, "Invalid value", verbose, json_output, start_time) + except Exception as e: + _handle_error(e, "Unexpected error", verbose, json_output, start_time) + + +def _handle_error( + e: Exception, + prefix: str, + verbose: bool, + json_output: bool, + start_time: float, +) -> None: + """Handle an error with appropriate output format.""" + import time as time_module + + duration = time_module.time() - start_time + error_info = get_error_info(e) + + if json_output: + json_data = format_error_json( + error_type=error_info.error_type, + message=f"{prefix}: {error_info.message}", + details=error_info.details, + recovery_suggestion=error_info.recovery_suggestion, + ) + json_data["meta"]["duration_seconds"] = duration + print_json(json_data) + else: + print_error(f"{prefix}: {error_info.message}") + if error_info.details: + console.print(f"[dim]{error_info.details}[/dim]") + if error_info.recovery_suggestion: + console.print(f"\n[info]Hint: {error_info.recovery_suggestion}[/info]") + if verbose and not isinstance(e, (ValueError, SQLInjectionError)): + console.print_exception() + + raise typer.Exit(2) from None diff --git a/src/quack_diff/cli/formatters.py b/src/quack_diff/cli/formatters.py index 0df38a7..9e36ec4 100644 --- a/src/quack_diff/cli/formatters.py +++ b/src/quack_diff/cli/formatters.py @@ -11,7 +11,7 @@ from rich.text import Text from quack_diff.cli.console import console -from quack_diff.core.differ import DiffType +from quack_diff.core.differ import CountResult, DiffType if TYPE_CHECKING: from quack_diff.core.differ import DiffResult, SchemaComparisonResult @@ -373,3 +373,49 @@ def print_snowflake_connections(connections: list[SnowflakeConnectionInfo]) -> N console.print() console.print(format_snowflake_connections(connections)) console.print() + + +def format_count_summary( + result: CountResult, + display_name_map: dict[str, str] | None = None, +) -> Panel: + """Create a summary panel for count check results. + + Args: + result: CountResult to format + display_name_map: Optional map from resolved table name to display name + + Returns: + Rich Panel with count table and status + """ + table = Table(show_header=True, box=None, padding=(0, 2)) + table.add_column("Table", style="bold") + table.add_column("Count", justify="right") + + display_name_map = display_name_map or {} + for tc in result.table_counts: + display_name = display_name_map.get(tc.table, tc.table) + table.add_row(display_name, f"{tc.count:,}") + + table.add_row("", "") # Spacer + status = Text("MATCH", style="success") if result.is_match else Text("MISMATCH", style="error") + table.add_row("Status", status) + + title = "Count Summary" + border_style = "green" if result.is_match else "red" + return Panel(table, title=title, border_style=border_style) + + +def print_count_result( + result: CountResult, + display_name_map: dict[str, str] | None = None, +) -> None: + """Print count check result to the console. + + Args: + result: CountResult to print + display_name_map: Optional map from resolved table name to display name + """ + console.print() + console.print(format_count_summary(result, display_name_map)) + console.print() diff --git a/src/quack_diff/cli/main.py b/src/quack_diff/cli/main.py index 39f4139..aa208aa 100644 --- a/src/quack_diff/cli/main.py +++ b/src/quack_diff/cli/main.py @@ -5,7 +5,7 @@ import typer from quack_diff import __version__ -from quack_diff.cli.commands import attach, compare, schema +from quack_diff.cli.commands import attach, compare, count, schema from quack_diff.cli.console import console app = typer.Typer( @@ -17,6 +17,7 @@ # Register commands app.command()(compare) +app.command()(count) app.command()(schema) app.command()(attach) diff --git a/src/quack_diff/cli/output.py b/src/quack_diff/cli/output.py index 80f9264..a6a6d5e 100644 --- a/src/quack_diff/cli/output.py +++ b/src/quack_diff/cli/output.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from quack_diff.core.differ import DiffResult, SchemaComparisonResult + from quack_diff.core.differ import CountResult, DiffResult, SchemaComparisonResult class OutputFormat(str, Enum): @@ -80,6 +80,19 @@ class JSONErrorOutput: error: dict[str, Any] = field(default_factory=dict) +@dataclass +class JSONCountOutput: + """JSON output structure for count command results.""" + + status: str # "match", "mismatch" + exit_code: int + meta: JSONOutputMeta + key_column: str | None + mode: str # "rows", "distinct" + tables: list[dict[str, Any]] + is_match: bool + + def get_version() -> str: """Get quack-diff version.""" try: @@ -226,6 +239,48 @@ def format_schema_result_json( return asdict(output) +def format_count_result_json( + result: CountResult, + display_name_map: dict[str, str] | None = None, + duration_seconds: float | None = None, +) -> dict[str, Any]: + """Format CountResult as JSON-serializable dictionary. + + Args: + result: CountResult to format + display_name_map: Optional map from resolved table name to display name + duration_seconds: Optional execution duration + + Returns: + Dictionary suitable for JSON serialization + """ + display_name_map = display_name_map or {} + status = "match" if result.is_match else "mismatch" + exit_code = 0 if result.is_match else 1 + + tables_data = [ + { + "table": display_name_map.get(tc.table, tc.table), + "count": tc.count, + } + for tc in result.table_counts + ] + + output = JSONCountOutput( + status=status, + exit_code=exit_code, + meta=JSONOutputMeta( + version=get_version(), + duration_seconds=duration_seconds, + ), + key_column=result.key_column, + mode=result.mode, + tables=tables_data, + is_match=result.is_match, + ) + return asdict(output) + + def format_attach_result_json( alias: str, path: str, diff --git a/src/quack_diff/core/differ.py b/src/quack_diff/core/differ.py index 0bc1312..e384d04 100644 --- a/src/quack_diff/core/differ.py +++ b/src/quack_diff/core/differ.py @@ -167,6 +167,35 @@ def is_match(self) -> bool: return self.total_differences == 0 +@dataclass +class TableCount: + """Row or distinct count for a single table.""" + + table: str + count: int + + +@dataclass +class CountResult: + """Result of comparing row counts across multiple tables.""" + + table_counts: list[TableCount] + key_column: str | None = None # None = COUNT(*), else COUNT(DISTINCT key) + is_match: bool = False + + @property + def mode(self) -> str: + """Count mode: 'rows' for COUNT(*), 'distinct' for COUNT(DISTINCT key).""" + return "distinct" if self.key_column else "rows" + + @property + def expected_count(self) -> int | None: + """Reference count (first table). None if no tables.""" + if not self.table_counts: + return None + return self.table_counts[0].count + + class DataDiffer: """Compares data between two database tables. @@ -379,6 +408,61 @@ def get_row_count( details="Verify the table name, schema, and database are correct", ) from e + def count_check( + self, + tables: list[str], + key_column: str | None = None, + dialect: Dialect | str = Dialect.DUCKDB, + ) -> CountResult: + """Check that all tables have the same row or distinct-key count. + + Args: + tables: List of fully qualified table names + key_column: If set, use COUNT(DISTINCT key_column); else COUNT(*) + dialect: SQL dialect for all tables + + Returns: + CountResult with per-table counts and is_match + + Raises: + TableNotFoundError: If any table does not exist + QueryExecutionError: If any count query fails + """ + if not tables: + return CountResult(table_counts=[], key_column=key_column, is_match=True) + + table_counts: list[TableCount] = [] + for table in tables: + if key_column: + query = self.query_builder.build_distinct_count_query( + table=table, + key_column=key_column, + dialect=dialect, + ) + else: + query = self.query_builder.build_count_query( + table=table, + dialect=dialect, + ) + try: + row = self.connector.execute_fetchone(query) + count = row[0] if row else 0 + except TableNotFoundError as e: + raise TableNotFoundError( + table=table, + message=f"Cannot count: table '{table}' does not exist", + details="Verify the table name, schema, and database are correct", + ) from e + table_counts.append(TableCount(table=table, count=count)) + + reference = table_counts[0].count + is_match = all(tc.count == reference for tc in table_counts) + return CountResult( + table_counts=table_counts, + key_column=key_column, + is_match=is_match, + ) + def _validate_key_column( self, key_column: str, diff --git a/src/quack_diff/core/query_builder.py b/src/quack_diff/core/query_builder.py index 4e6b4ad..d9d2679 100644 --- a/src/quack_diff/core/query_builder.py +++ b/src/quack_diff/core/query_builder.py @@ -181,6 +181,44 @@ def build_count_query( return f"SELECT COUNT(*) AS row_count FROM {table_ref}" + def build_distinct_count_query( + self, + table: str, + key_column: str, + dialect: Dialect | str = Dialect.DUCKDB, + timestamp: str | None = None, + offset: str | None = None, + ) -> str: + """Build a query that counts distinct values in a key column. + + Args: + table: Fully qualified table name + key_column: Column to count distinct values for + dialect: SQL dialect + timestamp: Optional timestamp for time-travel + offset: Optional offset for time-travel + + Returns: + SQL query string + + Raises: + SQLInjectionError: If table or column name contains unsafe characters + """ + adapter = self.get_adapter(dialect) + + sanitized_table = sanitize_identifier(table) + sanitized_key = sanitize_identifier(key_column) + + table_ref = sanitized_table + if timestamp is not None or offset is not None: + table_ref = adapter.wrap_table_with_time_travel( + table=sanitized_table, + timestamp=timestamp, + offset=offset, + ) + + return f"SELECT COUNT(DISTINCT {sanitized_key}) AS row_count FROM {table_ref}" + def build_schema_query( self, table: str, diff --git a/tests/cli/test_main.py b/tests/cli/test_main.py index 912692b..6bd2845 100644 --- a/tests/cli/test_main.py +++ b/tests/cli/test_main.py @@ -42,6 +42,7 @@ def test_help_shows_available_commands(self): """Test that --help lists available commands.""" result = runner.invoke(app, ["--help"]) assert "compare" in result.output + assert "count" in result.output assert "schema" in result.output assert "attach" in result.output @@ -82,3 +83,80 @@ def test_no_args_shows_help(self): # Typer/Click exits with code 2 when showing help due to missing args assert result.exit_code == 2 assert "quack-diff" in result.output + + +class TestCountCommand: + """Test count command.""" + + def test_count_help(self): + """Test count --help shows options.""" + result = runner.invoke(app, ["count", "--help"]) + assert result.exit_code == 0 + assert "count" in result.output.lower() + assert "--tables" in result.output or "-t" in result.output + assert "--key" in result.output or "-k" in result.output + assert "COUNT" in result.output or "count" in result.output + + def test_count_requires_at_least_two_tables(self): + """Test count with one table fails with clear message.""" + result = runner.invoke(app, ["count", "-t", "some_table"]) + assert result.exit_code == 2 + assert "two tables" in result.output.lower() or "At least two" in result.output + + def test_count_with_two_tables_match(self, temp_duckdb_count_tables): + """Test count with two tables (same count) exits 0.""" + config_path, table_refs = temp_duckdb_count_tables + result = runner.invoke( + app, + ["count", "-t", table_refs[0], "-t", table_refs[1], "--config", config_path], + ) + assert result.exit_code == 0 + assert "match" in result.output.lower() or "Match" in result.output + + def test_count_with_key_distinct(self, temp_duckdb_count_tables): + """Test count --key with two matching tables.""" + config_path, table_refs = temp_duckdb_count_tables + result = runner.invoke( + app, + [ + "count", + "-t", + table_refs[0], + "-t", + table_refs[1], + "--key", + "id", + "--config", + config_path, + ], + ) + assert result.exit_code == 0 + assert "match" in result.output.lower() or "Match" in result.output + + def test_count_json_output(self, temp_duckdb_count_tables): + """Test count --json produces valid JSON with expected structure.""" + import json + + config_path, table_refs = temp_duckdb_count_tables + result = runner.invoke( + app, + [ + "count", + "-t", + table_refs[0], + "-t", + table_refs[1], + "--json", + "--config", + config_path, + ], + ) + assert result.exit_code == 0 + data = json.loads(result.output) + assert data["status"] == "match" + assert data["exit_code"] == 0 + assert "tables" in data + assert len(data["tables"]) == 2 + assert data["is_match"] is True + assert "meta" in data + assert data["mode"] in ("rows", "distinct") diff --git a/tests/conftest.py b/tests/conftest.py index 7a96057..3249df5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,9 @@ from __future__ import annotations +import tempfile +from pathlib import Path + import duckdb import pytest @@ -170,3 +173,46 @@ def schema_mismatch_tables(connector: DuckDBConnector): """) return connector + + +@pytest.fixture +def temp_parquet_files(): + """Create temporary parquet files for CLI testing (e.g. count command).""" + with tempfile.TemporaryDirectory() as tmpdir: + source_path = Path(tmpdir) / "source.parquet" + target_path = Path(tmpdir) / "target.parquet" + + conn = duckdb.connect() + conn.execute("CREATE TABLE source (id INT, name VARCHAR, value INT)") + conn.execute("INSERT INTO source VALUES (1, 'a', 100), (2, 'b', 200)") + conn.execute(f"COPY source TO '{source_path}'") + + conn.execute("CREATE TABLE target (id INT, name VARCHAR, value INT)") + conn.execute("INSERT INTO target VALUES (1, 'a', 100), (2, 'b', 200)") + conn.execute(f"COPY target TO '{target_path}'") + conn.close() + + yield str(source_path), str(target_path) + + +@pytest.fixture +def temp_duckdb_count_tables(): + """Create a temporary DuckDB file with two tables (same row count) and config for count CLI tests.""" + with tempfile.TemporaryDirectory() as tmpdir: + tmpdir_path = Path(tmpdir) + db_path = tmpdir_path / "count_test.duckdb" + config_path = tmpdir_path / "quack-diff.yaml" + + conn = duckdb.connect(str(db_path)) + conn.execute("CREATE TABLE t1 (id INT, name VARCHAR)") + conn.execute("INSERT INTO t1 VALUES (1, 'a'), (2, 'b'), (3, 'c')") + conn.execute("CREATE TABLE t2 (id INT, name VARCHAR)") + conn.execute("INSERT INTO t2 VALUES (1, 'a'), (2, 'b'), (3, 'c')") + conn.close() + + config_path.write_text( + f"databases:\n countdb:\n type: duckdb\n path: {db_path}\n", + encoding="utf-8", + ) + + yield str(config_path), ["countdb.t1", "countdb.t2"] diff --git a/tests/core/test_differ.py b/tests/core/test_differ.py index ece8b53..63dc911 100644 --- a/tests/core/test_differ.py +++ b/tests/core/test_differ.py @@ -5,10 +5,12 @@ from quack_diff.core.connector import DuckDBConnector from quack_diff.core.differ import ( ColumnInfo, + CountResult, DataDiffer, DiffResult, DiffType, SchemaComparisonResult, + TableCount, ) from quack_diff.core.sql_utils import KeyColumnError, SchemaError, TableNotFoundError @@ -182,6 +184,127 @@ def test_quick_check_different_tables(self, sample_tables: DuckDBConnector): assert is_match is False + def test_count_check_matching_tables(self, identical_tables: DuckDBConnector): + """Test count_check when all tables have the same row count.""" + differ = DataDiffer(identical_tables) + result = differ.count_check( + tables=["identical_source", "identical_target"], + key_column=None, + ) + + assert isinstance(result, CountResult) + assert result.is_match + assert len(result.table_counts) == 2 + assert result.table_counts[0].count == result.table_counts[1].count == 3 + assert result.mode == "rows" + assert result.expected_count == 3 + + def test_count_check_mismatching_tables(self, connector: DuckDBConnector): + """Test count_check when tables have different row counts.""" + connector.execute("CREATE TABLE small (id INT)") + connector.execute("INSERT INTO small VALUES (1), (2)") + connector.execute("CREATE TABLE large (id INT)") + connector.execute("INSERT INTO large VALUES (1), (2), (3), (4)") + + differ = DataDiffer(connector) + result = differ.count_check( + tables=["small", "large"], + key_column=None, + ) + + assert isinstance(result, CountResult) + assert not result.is_match + assert result.table_counts[0].table == "small" + assert result.table_counts[0].count == 2 + assert result.table_counts[1].table == "large" + assert result.table_counts[1].count == 4 + + def test_count_check_three_tables_mismatch(self, connector: DuckDBConnector): + """Test count_check with three tables where one has different count.""" + connector.execute("CREATE TABLE t1 (id INT)") + connector.execute("INSERT INTO t1 VALUES (1), (2), (3)") + connector.execute("CREATE TABLE t2 (id INT)") + connector.execute("INSERT INTO t2 VALUES (1), (2), (3)") + connector.execute("CREATE TABLE t3 (id INT)") + connector.execute("INSERT INTO t3 VALUES (1), (2)") + + differ = DataDiffer(connector) + result = differ.count_check(tables=["t1", "t2", "t3"], key_column=None) + + assert not result.is_match + assert [tc.count for tc in result.table_counts] == [3, 3, 2] + + def test_count_check_distinct_key(self, identical_tables: DuckDBConnector): + """Test count_check with COUNT(DISTINCT key).""" + differ = DataDiffer(identical_tables) + result = differ.count_check( + tables=["identical_source", "identical_target"], + key_column="id", + ) + + assert result.is_match + assert result.mode == "distinct" + assert result.key_column == "id" + assert result.table_counts[0].count == 3 + assert result.table_counts[1].count == 3 + + def test_count_check_empty_tables_list(self, connector: DuckDBConnector): + """Test count_check with empty table list returns match.""" + differ = DataDiffer(connector) + result = differ.count_check(tables=[], key_column=None) + + assert result.is_match + assert result.table_counts == [] + assert result.expected_count is None + + def test_count_check_single_table(self, identical_tables: DuckDBConnector): + """Test count_check with single table is match.""" + differ = DataDiffer(identical_tables) + result = differ.count_check( + tables=["identical_source"], + key_column=None, + ) + + assert result.is_match + assert len(result.table_counts) == 1 + assert result.table_counts[0].count == 3 + + +class TestCountResult: + """Tests for CountResult dataclass.""" + + def test_mode_rows_when_no_key(self): + """Test mode is 'rows' when key_column is None.""" + result = CountResult( + table_counts=[TableCount("a", 10), TableCount("b", 10)], + key_column=None, + is_match=True, + ) + assert result.mode == "rows" + + def test_mode_distinct_when_key_set(self): + """Test mode is 'distinct' when key_column is set.""" + result = CountResult( + table_counts=[TableCount("a", 10), TableCount("b", 10)], + key_column="id", + is_match=True, + ) + assert result.mode == "distinct" + + def test_expected_count_first_table(self): + """Test expected_count is first table's count.""" + result = CountResult( + table_counts=[TableCount("a", 100), TableCount("b", 100)], + key_column=None, + is_match=True, + ) + assert result.expected_count == 100 + + def test_expected_count_empty(self): + """Test expected_count is None when no tables.""" + result = CountResult(table_counts=[], key_column=None, is_match=True) + assert result.expected_count is None + class TestDiffResult: """Tests for DiffResult dataclass.""" diff --git a/tests/core/test_query_builder.py b/tests/core/test_query_builder.py index e617586..5307e5e 100644 --- a/tests/core/test_query_builder.py +++ b/tests/core/test_query_builder.py @@ -53,6 +53,29 @@ def test_build_count_query(self, builder: QueryBuilder): assert "SELECT COUNT(*)" in query assert "my_table" in query + def test_build_distinct_count_query(self, builder: QueryBuilder): + """Test distinct count query generation.""" + query = builder.build_distinct_count_query( + table="my_table", + key_column="id", + dialect=Dialect.DUCKDB, + ) + + assert "SELECT COUNT(DISTINCT" in query + assert "id" in query + assert "my_table" in query + assert "row_count" in query + + def test_build_distinct_count_query_snowflake(self, builder: QueryBuilder): + """Test distinct count query for Snowflake dialect.""" + query = builder.build_distinct_count_query( + table="sf.schema.t", + key_column="order_id", + dialect=Dialect.DUCKDB, + ) + assert "COUNT(DISTINCT" in query + assert "order_id" in query + def test_build_schema_query(self, builder: QueryBuilder): """Test schema query generation.""" query = builder.build_schema_query(