feat: Implement connection resolution and materialization for SQL inputs; enhance concurrency handling in stats checks

luisggoncalves · luisggoncalves · commit 57664611cb48 · 2026-04-02T10:24:40.000-03:00
diff --git a/sqlcompare/helpers.py b/sqlcompare/helpers.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import os
 import re
 from dataclasses import dataclass
 from pathlib import Path
@@ -45,6 +46,66 @@ def detect_input(value: str) -> InputSpec:
 
     return InputSpec(kind="table", value=value.strip(), source="none")
 
+
+def resolve_connection(
+    connection: str | None,
+    *,
+    error_cls: type[Exception] = typer.BadParameter,
+) -> str:
+    """Resolve a connection from an explicit value or the default environment variables."""
+    if connection:
+        return connection
+
+    default_conn = os.getenv("SQLCOMPARE_CONN_DEFAULT") or os.getenv("DTK_CONN_DEFAULT")
+    if not default_conn:
+        raise error_cls(
+            "No connection specified. Use --connection or set SQLCOMPARE_CONN_DEFAULT."
+        )
+    return default_conn
+
+
+def resolve_materialized_tables(
+    previous_spec: InputSpec,
+    current_spec: InputSpec,
+    *,
+    schema: str | None,
+    prefix: str,
+    suffix: str,
+) -> tuple[str, str]:
+    """Return table names for a pair of table-or-SQL inputs."""
+    schema_prefix = f"{schema}." if schema else ""
+    previous_name = (
+        previous_spec.value
+        if previous_spec.kind == "table"
+        else f"{schema_prefix}{prefix}_{suffix}_previous"
+    )
+    current_name = (
+        current_spec.value
+        if current_spec.kind == "table"
+        else f"{schema_prefix}{prefix}_{suffix}_new"
+    )
+    return previous_name, current_name
+
+
+def materialize_sql_inputs(
+    db: DBConnection,
+    *,
+    previous_spec: InputSpec,
+    current_spec: InputSpec,
+    previous_table: str,
+    current_table: str,
+    schema: str | None,
+) -> None:
+    """Create tables for SQL inputs while leaving table-name inputs untouched."""
+    if previous_spec.kind != "sql" and current_spec.kind != "sql":
+        return
+
+    ensure_schema(db, schema or "")
+    if previous_spec.kind == "sql":
+        create_table_from_select(db, previous_table, previous_spec.value)
+    if current_spec.kind == "sql":
+        create_table_from_select(db, current_table, current_spec.value)
+
 def expand_dataset_value(value: Any, base_dir: Path) -> Any:
     """
     Recursively expand template variables in dataset configuration values.
diff --git a/sqlcompare/run.py b/sqlcompare/run.py
@@ -5,11 +5,11 @@
 
 import typer
 
-from sqlcompare.helpers import detect_input
+from sqlcompare.helpers import detect_input, resolve_connection
 from sqlcompare.run_table import compare_table, run_table_cmd
 from sqlcompare.run_dataset import run_dataset_cmd
 from sqlcompare.run_stats import run_stats_cmd
-from sqlcompare.run_auto import _resolve_connection, run_auto_cmd
+from sqlcompare.run_auto import run_auto_cmd
 from sqlcompare.config import get_default_schema
 from sqlcompare.db import DBConnection
 from sqlcompare.helpers import create_table_from_select, ensure_schema
@@ -91,7 +91,7 @@ def run_query_cmd(
     schema = schema or get_default_schema()
     previous_sql = _load_sql_input(previous, "previous")
     current_sql = _load_sql_input(current, "current")
-    connection = _resolve_connection(connection)
+    connection = resolve_connection(connection)
 
     schema_prefix = f"{schema}." if schema else ""
     suffix = _suffix_from_paths(previous, current)
diff --git a/sqlcompare/run_auto.py b/sqlcompare/run_auto.py
@@ -1,27 +1,20 @@
 from __future__ import annotations
 
-import os
 import uuid
 
 import typer
 
 from sqlcompare.config import get_default_schema
 from sqlcompare.db import DBConnection
-from sqlcompare.helpers import create_table_from_select, detect_input, ensure_schema
+from sqlcompare.helpers import (
+    detect_input,
+    materialize_sql_inputs,
+    resolve_connection,
+    resolve_materialized_tables,
+)
 from sqlcompare.run_table import compare_table
 
 
-def _resolve_connection(connection: str | None) -> str:
-    if connection:
-        return connection
-    default_conn = os.getenv("SQLCOMPARE_CONN_DEFAULT") or os.getenv("DTK_CONN_DEFAULT")
-    if not default_conn:
-        raise typer.BadParameter(
-            "No connection specified. Use --connection or set SQLCOMPARE_CONN_DEFAULT."
-        )
-    return default_conn
-
-
 def run_auto_cmd(
     previous: str = typer.Argument(
         ..., help="Previous table name, CSV/XLSX file path, or SQL"
@@ -90,26 +83,24 @@ def run_auto_cmd(
         return
 
     if prev_spec.kind == "sql" or new_spec.kind == "sql":
-        connection = _resolve_connection(connection)
-        schema_prefix = f"{schema}." if schema else ""
-        suffix = uuid.uuid4().hex[:8]
-        previous_table = (
-            prev_spec.value
-            if prev_spec.kind == "table"
-            else f"{schema_prefix}sqlcompare_sql_{suffix}_previous"
-        )
-        new_table = (
-            new_spec.value
-            if new_spec.kind == "table"
-            else f"{schema_prefix}sqlcompare_sql_{suffix}_new"
+        connection = resolve_connection(connection)
+        previous_table, new_table = resolve_materialized_tables(
+            prev_spec,
+            new_spec,
+            schema=schema,
+            prefix="sqlcompare_sql",
+            suffix=uuid.uuid4().hex[:8],
         )
 
         with DBConnection(connection) as db:
-            ensure_schema(db, schema)
-            if prev_spec.kind == "sql":
-                create_table_from_select(db, previous_table, prev_spec.value)
-            if new_spec.kind == "sql":
-                create_table_from_select(db, new_table, new_spec.value)
+            materialize_sql_inputs(
+                db,
+                previous_spec=prev_spec,
+                current_spec=new_spec,
+                previous_table=previous_table,
+                current_table=new_table,
+                schema=schema,
+            )
 
         compare_table(
             previous_table,
diff --git a/sqlcompare/stats/runner.py b/sqlcompare/stats/runner.py
@@ -1,29 +1,24 @@
 from __future__ import annotations
 
-import os
+from dataclasses import replace
 import uuid
 from pathlib import Path
 
 import typer
 
 from sqlcompare.config import get_default_schema
 from sqlcompare.db import DBConnection
-from sqlcompare.helpers import create_table_from_select, detect_input, ensure_schema
+from sqlcompare.helpers import (
+    detect_input,
+    materialize_sql_inputs,
+    resolve_connection,
+    resolve_materialized_tables,
+)
 from sqlcompare.log import log
 from sqlcompare.stats.checks import get_check_map
 from sqlcompare.stats.models import ColumnPair, StatsContext
 from sqlcompare.stats.report import render_report
-
-
-def _resolve_connection(connection: str | None) -> str:
-    if connection:
-        return connection
-    default_conn = os.getenv("SQLCOMPARE_CONN_DEFAULT") or os.getenv("DTK_CONN_DEFAULT")
-    if not default_conn:
-        raise ValueError(
-            "No connection specified. Use --connection or set SQLCOMPARE_CONN_DEFAULT."
-        )
-    return default_conn
+from sqlcompare.utils.concurrency import run_ordered
 
 
 def _resolve_checks(checks: str | None) -> list[str]:
@@ -101,6 +96,29 @@ def _build_context(db: DBConnection, previous_name: str, current_name: str) -> S
     )
 
 
+def _run_selected_checks(
+    context: StatsContext,
+    selected_check_names: list[str],
+    check_map: dict[str, object],
+    connection_id: str | None,
+    parallel_safe: bool,
+) -> list:
+    def run_check(name: str):
+        definition = check_map[name]
+        if not parallel_safe:
+            return definition.runner(context, definition)
+
+        with DBConnection(connection_id) as db:
+            return definition.runner(replace(context, db=db), definition)
+
+    return run_ordered(
+        selected_check_names,
+        run_check,
+        enabled=parallel_safe,
+        max_workers=4,
+    )
+
+
 def compare_table_stats(
     table1: str,
     table2: str,
@@ -125,42 +143,63 @@ def compare_table_stats(
         table1_name = Path(spec_prev.value).stem
         table2_name = Path(spec_new.value).stem
     elif spec_prev.kind == "sql" or spec_new.kind == "sql":
-        connection_id = _resolve_connection(connection)
+        connection_id = resolve_connection(connection, error_cls=ValueError)
         schema = get_default_schema()
-        schema_prefix = f"{schema}." if schema else ""
         suffix = uuid.uuid4().hex[:8]
-        table1_name = (
-            spec_prev.value
-            if spec_prev.kind == "table"
-            else f"{schema_prefix}sqlcompare_stats_{suffix}_previous"
-        )
-        table2_name = (
-            spec_new.value
-            if spec_new.kind == "table"
-            else f"{schema_prefix}sqlcompare_stats_{suffix}_new"
+        table1_name, table2_name = resolve_materialized_tables(
+            spec_prev,
+            spec_new,
+            schema=schema,
+            prefix="sqlcompare_stats",
+            suffix=suffix,
         )
     else:
         table1_name = spec_prev.value
         table2_name = spec_new.value
 
-    with DBConnection(connection_id) as db:
+    parallel_safe = not (
+        spec_prev.kind == "file" and spec_new.kind == "file" and connection is None
+    )
+
+    def prepare_inputs(db: DBConnection) -> None:
         if spec_prev.kind == "file" and spec_new.kind == "file":
             if connection is None and connection_id == "duckdb:///:memory:":
                 db.create_table_from_file(table1_name, spec_prev.value)
                 db.create_table_from_file(table2_name, spec_new.value)
-        if spec_prev.kind == "sql" or spec_new.kind == "sql":
-            schema = get_default_schema()
-            ensure_schema(db, schema)
-            if spec_prev.kind == "sql":
-                create_table_from_select(db, table1_name, spec_prev.value)
-            if spec_new.kind == "sql":
-                create_table_from_select(db, table2_name, spec_new.value)
-
-        context = _build_context(db, table1_name, table2_name)
-        check_map = get_check_map()
-        results = [
-            check_map[name].runner(context, check_map[name])
-            for name in selected_check_names
-        ]
+        materialize_sql_inputs(
+            db,
+            previous_spec=spec_prev,
+            current_spec=spec_new,
+            previous_table=table1_name,
+            current_table=table2_name,
+            schema=get_default_schema(),
+        )
+
+    if not parallel_safe:
+        with DBConnection(connection_id) as db:
+            prepare_inputs(db)
+            context = _build_context(db, table1_name, table2_name)
+            check_map = get_check_map()
+            results = _run_selected_checks(
+                context=context,
+                selected_check_names=selected_check_names,
+                check_map=check_map,
+                connection_id=connection_id,
+                parallel_safe=False,
+            )
+    else:
+        with DBConnection(connection_id) as db:
+            prepare_inputs(db)
+
+        with DBConnection(connection_id) as db:
+            context = _build_context(db, table1_name, table2_name)
+            check_map = get_check_map()
+            results = _run_selected_checks(
+                context=context,
+                selected_check_names=selected_check_names,
+                check_map=check_map,
+                connection_id=connection_id,
+                parallel_safe=True,
+            )
 
     log.info(render_report(context, selected_check_names, results))
diff --git a/sqlcompare/utils/concurrency.py b/sqlcompare/utils/concurrency.py
@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+from typing import Callable, Sequence, TypeVar
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+def run_ordered(
+    items: Sequence[T],
+    worker: Callable[[T], R],
+    *,
+    enabled: bool = True,
+    max_workers: int | None = None,
+) -> list[R]:
+    """Run work items concurrently while preserving the input order in the results."""
+    if len(items) <= 1 or not enabled:
+        return [worker(item) for item in items]
+
+    worker_count = max_workers or len(items)
+    with ThreadPoolExecutor(max_workers=min(len(items), worker_count)) as executor:
+        futures = [executor.submit(worker, item) for item in items]
+        return [future.result() for future in futures]
diff --git a/tests/test_stats_comparison.py b/tests/test_stats_comparison.py
