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

data_diff/queries/base.py

@@ -0,0 +1,24 @@
+from typing import Generator
+
+import attrs
+
+
+@attrs.define(frozen=True)
+class _SKIP:
+    def __repr__(self) -> str:
+        return "SKIP"
+
+
+SKIP = _SKIP()
+
+
+class SqeletonError(Exception):
+    pass
+
+
+def args_as_tuple(exprs):
+    if len(exprs) == 1:
+        (e,) = exprs
+        if isinstance(e, Generator):
+            return tuple(e)
+    return exprs

data_diff/queries/extras.py

@@ -0,0 +1,29 @@
+"Useful AST classes that don't quite fall within the scope of regular SQL"
+
+from typing import Callable, Optional, Sequence
+
+import attrs
+
+from data_diff.abcs.database_types import ColType
+from data_diff.queries.ast_classes import Expr, ExprNode
+
+
+@attrs.define(frozen=True)
+class NormalizeAsString(ExprNode):
+    expr: ExprNode
+    expr_type: Optional[ColType] = None
+
+    @property
+    def type(self) -> Optional[type]:
+        return str
+
+
+@attrs.define(frozen=True)
+class ApplyFuncAndNormalizeAsString(ExprNode):
+    expr: ExprNode
+    apply_func: Optional[Callable] = None
+
+
+@attrs.define(frozen=True)
+class Checksum(ExprNode):
+    exprs: Sequence[Expr]

data_diff/query_utils.py

@@ -0,0 +1,56 @@
+"Module for query utilities that didn't make it into the query-builder (yet)"
+
+from contextlib import suppress
+
+from data_diff.abcs.database_types import DbPath
+from data_diff.databases.base import QueryError
+from data_diff.databases.oracle import Oracle
+from data_diff.queries.api import table, commit, Expr
+
+
+def _drop_table_oracle(name: DbPath):
+    t = table(name)
+    # Experience shows double drop is necessary
+    with suppress(QueryError):
+        yield t.drop()
+        yield t.drop()
+    yield commit
+
+
+def _drop_table(name: DbPath):
+    t = table(name)
+    yield t.drop(if_exists=True)
+    yield commit
+
+
+def drop_table(db, tbl) -> None:
+    if isinstance(db, Oracle):
+        db.query(_drop_table_oracle(tbl))
+    else:
+        db.query(_drop_table(tbl))
+
+
+def _append_to_table_oracle(path: DbPath, expr: Expr):
+    """See append_to_table"""
+    assert expr.schema, expr
+    t = table(path, schema=expr.schema)
+    with suppress(QueryError):
+        yield t.create()  # uses expr.schema
+        yield commit
+    yield t.insert_expr(expr)
+    yield commit
+
+
+def _append_to_table(path: DbPath, expr: Expr):
+    """Append to table"""
+    assert expr.schema, expr
+    t = table(path, schema=expr.schema)
+    yield t.create(if_not_exists=True)  # uses expr.schema
+    yield commit
+    yield t.insert_expr(expr)
+    yield commit
+
+
+def append_to_table(db, path, expr) -> None:
+    f = _append_to_table_oracle if isinstance(db, Oracle) else _append_to_table
+    db.query(f(path, expr))

data_diff/schema.py

@@ -0,0 +1,52 @@
+import logging
+from typing import Any, Collection, Iterator, Optional
+
+import attrs
+
+from data_diff.utils import CaseAwareMapping, CaseInsensitiveDict, CaseSensitiveDict
+from data_diff.abcs.database_types import DbPath
+
+logger = logging.getLogger("schema")
+
+Schema = CaseAwareMapping
+
+
+@attrs.frozen(kw_only=True)
+class RawColumnInfo(Collection[Any]):
+    """
+    A raw row representing the schema info about a column.
+
+    Do not rely on this class too much, it will be removed soon when the schema
+    selecting & parsing methods are united into one overrideable method.
+    """
+
+    column_name: str
+    data_type: str
+    datetime_precision: Optional[int] = None
+    numeric_precision: Optional[int] = None
+    numeric_scale: Optional[int] = None
+    collation_name: Optional[str] = None
+
+    # It was a tuple once, so we keep it backward compatible temporarily, until remade to classes.
+    def __iter__(self) -> Iterator[Any]:
+        return iter(
+            (self.column_name, self.data_type, self.datetime_precision, self.numeric_precision, self.numeric_scale)
+        )
+
+    def __len__(self) -> int:
+        return 5
+
+    def __contains__(self, item: Any) -> bool:
+        return False  # that was not used
+
+
+def create_schema(db_name: str, table_path: DbPath, schema: dict, case_sensitive: bool) -> CaseAwareMapping:
+    logger.info(f"[{db_name}] Schema = {schema}")
+
+    if case_sensitive:
+        return CaseSensitiveDict(schema)
+
+    if len({k.lower() for k in schema}) < len(schema):
+        logger.warning(f'Ambiguous schema for {db_name}:{".".join(table_path)} | Columns = {", ".join(list(schema))}')
+        logger.warning("We recommend to disable case-insensitivity (set --case-sensitive).")
+    return CaseInsensitiveDict(schema)

data_diff/table_segment.py

@@ -0,0 +1,286 @@
+import time
+from typing import Container, Dict, List, Optional, Sequence, Tuple
+import logging
+from itertools import product
+
+import attrs
+from typing_extensions import Self
+
+from data_diff.utils import safezip, Vector
+from data_diff.utils import ArithString, split_space
+from data_diff.databases.base import Database
+from data_diff.abcs.database_types import DbPath, DbKey, DbTime, IKey
+from data_diff.schema import RawColumnInfo, Schema, create_schema
+from data_diff.queries.extras import Checksum
+from data_diff.queries.api import Count, SKIP, table, this, Expr, min_, max_, Code
+from data_diff.queries.extras import ApplyFuncAndNormalizeAsString, NormalizeAsString
+
+logger = logging.getLogger("table_segment")
+
+RECOMMENDED_CHECKSUM_DURATION = 20
+
+
+def split_key_space(min_key: DbKey, max_key: DbKey, count: int) -> List[DbKey]:
+    assert min_key < max_key
+
+    if max_key - min_key <= count:
+        count = 1
+
+    if isinstance(min_key, ArithString):
+        assert type(min_key) is type(max_key)
+        checkpoints = min_key.range(max_key, count)
+    else:
+        checkpoints = split_space(min_key, max_key, count)
+
+    assert all(min_key < x < max_key for x in checkpoints)
+    return [min_key] + checkpoints + [max_key]
+
+
+def int_product(nums: List[int]) -> int:
+    p = 1
+    for n in nums:
+        p *= n
+    return p
+
+
+def split_compound_key_space(mn: Vector, mx: Vector, count: int) -> List[List[DbKey]]:
+    """Returns a list of split-points for each key dimension, essentially returning an N-dimensional grid of split points."""
+    return [split_key_space(mn_k, mx_k, count) for mn_k, mx_k in safezip(mn, mx)]
+
+
+def create_mesh_from_points(*values_per_dim: list) -> List[Tuple[Vector, Vector]]:
+    """Given a list of values along each axis of N dimensional space,
+    return an array of boxes whose start-points & end-points align with the given values,
+    and together consitute a mesh filling that space entirely (within the bounds of the given values).
+
+    Assumes given values are already ordered ascending.
+
+    len(boxes) == ∏i( len(i)-1 )
+
+    Example:
+        ::
+            >>> d1 = 'a', 'b', 'c'
+            >>> d2 = 1, 2, 3
+            >>> d3 = 'X', 'Y'
+            >>> create_mesh_from_points(d1, d2, d3)
+            [
+                [('a', 1, 'X'), ('b', 2, 'Y')],
+                [('a', 2, 'X'), ('b', 3, 'Y')],
+                [('b', 1, 'X'), ('c', 2, 'Y')],
+                [('b', 2, 'X'), ('c', 3, 'Y')]
+            ]
+    """
+    assert all(len(v) >= 2 for v in values_per_dim), values_per_dim
+
+    # Create tuples of (v1, v2) for each pair of adjacent values
+    ranges = [list(zip(values[:-1], values[1:])) for values in values_per_dim]
+
+    assert all(a <= b for r in ranges for a, b in r)
+
+    # Create a product of all the ranges
+    res = [tuple(Vector(a) for a in safezip(*r)) for r in product(*ranges)]
+
+    expected_len = int_product(len(v) - 1 for v in values_per_dim)
+    assert len(res) == expected_len, (len(res), expected_len)
+    return res
+
+
+@attrs.define(frozen=True)
+class TableSegment:
+    """Signifies a segment of rows (and selected columns) within a table
+
+    Parameters:
+        database (Database): Database instance. See :meth:`connect`
+        table_path (:data:`DbPath`): Path to table in form of a tuple. e.g. `('my_dataset', 'table_name')`
+        key_columns (Tuple[str]): Name of the key column, which uniquely identifies each row (usually id)
+        update_column (str, optional): Name of updated column, which signals that rows changed.
+                                       Usually updated_at or last_update. Used by `min_update` and `max_update`.
+        extra_columns (Tuple[str, ...], optional): Extra columns to compare
+        min_key (:data:`Vector`, optional): Lowest key value, used to restrict the segment
+        max_key (:data:`Vector`, optional): Highest key value, used to restrict the segment
+        min_update (:data:`DbTime`, optional): Lowest update_column value, used to restrict the segment
+        max_update (:data:`DbTime`, optional): Highest update_column value, used to restrict the segment
+        where (str, optional): An additional 'where' expression to restrict the search space.
+
+        case_sensitive (bool): If false, the case of column names will adjust according to the schema. Default is true.
+
+    """
+
+    # Location of table
+    database: Database
+    table_path: DbPath
+
+    # Columns
+    key_columns: Tuple[str, ...]
+    update_column: Optional[str] = None
+    extra_columns: Tuple[str, ...] = ()
+    ignored_columns: Container[str] = frozenset()
+
+    # Restrict the segment
+    min_key: Optional[Vector] = None
+    max_key: Optional[Vector] = None
+    min_update: Optional[DbTime] = None
+    max_update: Optional[DbTime] = None
+    where: Optional[str] = None
+
+    case_sensitive: Optional[bool] = True
+    _schema: Optional[Schema] = None
+
+    def __attrs_post_init__(self) -> None:
+        if not self.update_column and (self.min_update or self.max_update):
+            raise ValueError("Error: the min_update/max_update feature requires 'update_column' to be set.")
+
+        if self.min_key is not None and self.max_key is not None and self.min_key >= self.max_key:
+            raise ValueError(f"Error: min_key expected to be smaller than max_key! ({self.min_key} >= {self.max_key})")
+
+        if self.min_update is not None and self.max_update is not None and self.min_update >= self.max_update:
+            raise ValueError(
+                f"Error: min_update expected to be smaller than max_update! ({self.min_update} >= {self.max_update})"
+            )
+
+    def _where(self) -> Optional[str]:
+        return f"({self.where})" if self.where else None
+
+    def _with_raw_schema(self, raw_schema: Dict[str, RawColumnInfo]) -> Self:
+        schema = self.database._process_table_schema(self.table_path, raw_schema, self.relevant_columns, self._where())
+        return self.new(schema=create_schema(self.database.name, self.table_path, schema, self.case_sensitive))
+
+    def with_schema(self) -> Self:
+        "Queries the table schema from the database, and returns a new instance of TableSegment, with a schema."
+        if self._schema:
+            return self
+
+        return self._with_raw_schema(self.database.query_table_schema(self.table_path))
+
+    def get_schema(self) -> Dict[str, RawColumnInfo]:
+        return self.database.query_table_schema(self.table_path)
+
+    def _make_key_range(self):
+        if self.min_key is not None:
+            for mn, k in safezip(self.min_key, self.key_columns):
+                yield mn <= this[k]
+        if self.max_key is not None:
+            for k, mx in safezip(self.key_columns, self.max_key):
+                yield this[k] < mx
+
+    def _make_update_range(self):
+        if self.min_update is not None:
+            yield self.min_update <= this[self.update_column]
+        if self.max_update is not None:
+            yield this[self.update_column] < self.max_update
+
+    @property
+    def source_table(self):
+        return table(*self.table_path, schema=self._schema)
+
+    def make_select(self):
+        return self.source_table.where(
+            *self._make_key_range(), *self._make_update_range(), Code(self._where()) if self.where else SKIP
+        )
+
+    def get_values(self) -> list:
+        "Download all the relevant values of the segment from the database"
+
+        # Fetch all the original columns, even if some were later excluded from checking.
+        fetched_cols = [NormalizeAsString(this[c]) for c in self.relevant_columns]
+        select = self.make_select().select(*fetched_cols)
+        return self.database.query(select, List[Tuple])
+
+    def choose_checkpoints(self, count: int) -> List[List[DbKey]]:
+        "Suggests a bunch of evenly-spaced checkpoints to split by, including start, end."
+
+        assert self.is_bounded
+
+        # Take Nth root of count, to approximate the appropriate box size
+        count = int(count ** (1 / len(self.key_columns))) or 1
+
+        return split_compound_key_space(self.min_key, self.max_key, count)
+
+    def segment_by_checkpoints(self, checkpoints: List[List[DbKey]]) -> List["TableSegment"]:
+        "Split the current TableSegment to a bunch of smaller ones, separated by the given checkpoints"
+
+        return [self.new_key_bounds(min_key=s, max_key=e) for s, e in create_mesh_from_points(*checkpoints)]
+
+    def new(self, **kwargs) -> Self:
+        """Creates a copy of the instance using 'replace()'"""
+        return attrs.evolve(self, **kwargs)
+
+    def new_key_bounds(self, min_key: Vector, max_key: Vector, *, key_types: Optional[Sequence[IKey]] = None) -> Self:
+        if self.min_key is not None:
+            assert self.min_key <= min_key, (self.min_key, min_key)
+            assert self.min_key < max_key
+
+        if self.max_key is not None:
+            assert min_key < self.max_key
+            assert max_key <= self.max_key
+
+        # If asked, enforce the PKs to proper types, mainly to meta-params of the relevant side,
+        # so that we do not leak e.g. casing of UUIDs from side A to side B and vice versa.
+        # If not asked, keep the meta-params of the keys as is (assume them already casted).
+        if key_types is not None:
+            min_key = Vector(type.make_value(val) for type, val in safezip(key_types, min_key))
+            max_key = Vector(type.make_value(val) for type, val in safezip(key_types, max_key))
+
+        return attrs.evolve(self, min_key=min_key, max_key=max_key)
+
+    @property
+    def relevant_columns(self) -> List[str]:
+        extras = list(self.extra_columns)
+
+        if self.update_column and self.update_column not in extras:
+            extras = [self.update_column] + extras
+
+        return list(self.key_columns) + extras
+
+    def count(self) -> int:
+        """Count how many rows are in the segment, in one pass."""
+        return self.database.query(self.make_select().select(Count()), int)
+
+    def count_and_checksum(self) -> Tuple[int, int]:
+        """Count and checksum the rows in the segment, in one pass."""
+
+        checked_columns = [c for c in self.relevant_columns if c not in self.ignored_columns]
+        cols = [NormalizeAsString(this[c]) for c in checked_columns]
+
+        start = time.monotonic()
+        q = self.make_select().select(Count(), Checksum(cols))
+        count, checksum = self.database.query(q, tuple)
+        duration = time.monotonic() - start
+        if duration > RECOMMENDED_CHECKSUM_DURATION:
+            logger.warning(
+                "Checksum is taking longer than expected (%.2f). "
+                "We recommend increasing --bisection-factor or decreasing --threads.",
+                duration,
+            )
+
+        if count:
+            assert checksum, (count, checksum)
+        return count or 0, int(checksum) if count else None
+
+    def query_key_range(self) -> Tuple[tuple, tuple]:
+        """Query database for minimum and maximum key. This is used for setting the initial bounds."""
+        # Normalizes the result (needed for UUIDs) after the min/max computation
+        select = self.make_select().select(
+            ApplyFuncAndNormalizeAsString(this[k], f) for k in self.key_columns for f in (min_, max_)
+        )
+        result = tuple(self.database.query(select, tuple))
+
+        if any(i is None for i in result):
+            raise ValueError("Table appears to be empty")
+
+        # Min/max keys are interleaved
+        min_key, max_key = result[::2], result[1::2]
+        assert len(min_key) == len(max_key)
+
+        return min_key, max_key
+
+    @property
+    def is_bounded(self):
+        return self.min_key is not None and self.max_key is not None
+
+    def approximate_size(self):
+        if not self.is_bounded:
+            raise RuntimeError("Cannot approximate the size of an unbounded segment. Must have min_key and max_key.")
+        diff = self.max_key - self.min_key
+        assert all(d > 0 for d in diff)
+        return int_product(diff)

data_diff/thread_utils.py

@@ -0,0 +1,98 @@
+import itertools
+from queue import PriorityQueue
+from collections import deque
+from collections.abc import Iterable
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures.thread import _WorkItem
+from time import sleep
+from typing import Any, Callable, Iterator, Optional
+
+import attrs
+
+
+class AutoPriorityQueue(PriorityQueue):
+    """Overrides PriorityQueue to automatically get the priority from _WorkItem.kwargs
+
+    We also assign a unique id for each item, to avoid making comparisons on _WorkItem.
+    As a side effect, items with the same priority are returned FIFO.
+    """
+
+    _counter = itertools.count().__next__
+
+    def put(self, item: Optional[_WorkItem], block=True, timeout=None) -> None:
+        priority = item.kwargs.pop("priority") if item is not None else 0
+        super().put((-priority, self._counter(), item), block, timeout)
+
+    def get(self, block=True, timeout=None) -> Optional[_WorkItem]:
+        _p, _c, work_item = super().get(block, timeout)
+        return work_item
+
+
+class PriorityThreadPoolExecutor(ThreadPoolExecutor):
+    """Overrides ThreadPoolExecutor to use AutoPriorityQueue
+
+    XXX WARNING: Might break in future versions of Python
+    """
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+        self._work_queue = AutoPriorityQueue()
+
+
+@attrs.define(frozen=False, init=False)
+class ThreadedYielder(Iterable):
+    """Yields results from multiple threads into a single iterator, ordered by priority.
+
+    To add a source iterator, call ``submit()`` with a function that returns an iterator.
+    Priority for the iterator can be provided via the keyword argument 'priority'. (higher runs first)
+    """
+
+    _pool: ThreadPoolExecutor
+    _futures: deque
+    _yield: deque
+    _exception: Optional[None]
+
+    _pool: ThreadPoolExecutor
+    _futures: deque
+    _yield: deque = attrs.field(alias="_yield")  # Python keyword!
+    _exception: Optional[None]
+    yield_list: bool
+
+    def __init__(self, max_workers: Optional[int] = None, yield_list: bool = False) -> None:
+        super().__init__()
+        self._pool = PriorityThreadPoolExecutor(max_workers)
+        self._futures = deque()
+        self._yield = deque()
+        self._exception = None
+        self.yield_list = yield_list
+
+    def _worker(self, fn, *args, **kwargs) -> None:
+        try:
+            res = fn(*args, **kwargs)
+            if res is not None:
+                if self.yield_list:
+                    self._yield.append(res)
+                else:
+                    self._yield += res
+        except Exception as e:
+            self._exception = e
+
+    def submit(self, fn: Callable, *args, priority: int = 0, **kwargs) -> None:
+        self._futures.append(self._pool.submit(self._worker, fn, *args, priority=priority, **kwargs))
+
+    def __iter__(self) -> Iterator[Any]:
+        while True:
+            if self._exception:
+                raise self._exception
+
+            while self._yield:
+                yield self._yield.popleft()
+
+            if not self._futures:
+                # No more tasks
+                return
+
+            if self._futures[0].done():
+                self._futures.popleft()
+            else:
+                sleep(0.001)