Merge pull request #165 from datafold/where

Added arbitrary 'where' option (issue #86)

Author: erezsh

README.md

@@ -205,6 +205,7 @@ Options:
   - `--bisection-factor` - Segments per iteration. When set to 2, it performs binary search.
   - `--bisection-threshold` - Minimal bisection threshold. i.e. maximum size of pages to diff locally.
   - `-j` or `--threads` - Number of worker threads to use per database. Default=1.
+  - `-w`, `--where` - An additional 'where' expression to restrict the search space.
   - `--conf`, `--run` - Specify the run and configuration from a TOML file. (see below)
 
 ### How to use with a configuration file

data_diff/__main__.py

@@ -71,6 +71,7 @@ def _remove_passwords_in_dict(d: dict):
     "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).",
 )
+@click.option("-w", "--where", default=None, help="An additional 'where' expression to restrict the search space.")
 @click.option(
     "--conf",
     default=None,
@@ -107,6 +108,7 @@ def _main(
     threads,
     keep_column_case,
     json_output,
+    where,
     threads1=None,
     threads2=None,
     __conf__=None,
@@ -166,6 +168,7 @@ def _main(
             min_update=max_age and parse_time_before_now(max_age),
             max_update=min_age and parse_time_before_now(min_age),
             case_sensitive=keep_column_case,
+            where=where,
         )
     except ParseError as e:
         logging.error("Error while parsing age expression: %s" % e)

data_diff/diff_tables.py

@@ -1,7 +1,6 @@
 """Provides classes for performing a table diff
 """
 
-from abc import ABC, abstractmethod
 import time
 import os
 from numbers import Number
@@ -13,7 +12,7 @@
 
 from runtype import dataclass
 
-from .sql import Select, Checksum, Compare, DbPath, DbKey, DbTime, Count, TableName, Time, Min, Max, Value
+from .sql import Select, Checksum, Compare, DbPath, DbKey, DbTime, Count, TableName, Time, Value
 from .utils import safezip, split_space
 from .databases.base import Database
 from .databases.database_types import (
@@ -50,6 +49,9 @@ class TableSegment:
         max_key (:data:`DbKey`, optional): Highest key_column 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.
 
     """
 
@@ -68,6 +70,7 @@ class TableSegment:
     min_update: DbTime = None
     max_update: DbTime = None
 
+    where: str = None
     case_sensitive: bool = True
     _schema: Schema = None
 
@@ -114,7 +117,7 @@ def _normalize_column(self, name: str, template: str = None) -> str:
         return self.database.normalize_value_by_type(col, col_type)
 
     def with_schema(self) -> "TableSegment":
-        "Queries the table schema from the database, and returns a new instance of TableSegmentWithSchema."
+        "Queries the table schema from the database, and returns a new instance of TableSegment, with a schema."
         if self._schema:
             return self
 
@@ -149,7 +152,12 @@ def _make_update_range(self):
     def _make_select(self, *, table=None, columns=None, where=None, group_by=None, order_by=None):
         if columns is None:
             columns = [self._normalize_column(self.key_column)]
-        where = list(self._make_key_range()) + list(self._make_update_range()) + ([] if where is None else [where])
+        where = [
+            *self._make_key_range(),
+            *self._make_update_range(),
+            *([] if where is None else [where]),
+            *([] if self.where is None else [self.where]),
+        ]
         order_by = None if order_by is None else [order_by]
         return Select(
             table=table or TableName(self.table_path),

tests/test_api.py

@@ -25,6 +25,9 @@ def setUp(self) -> None:
         )
         self.now = now = arrow.get(self.preql.now())
         self.preql.add(now, "now")
+        self.preql.add(now, self.now.shift(seconds=-10))
+        self.preql.add(now, self.now.shift(seconds=-7))
+        self.preql.add(now, self.now.shift(seconds=-6))
 
         self.preql(
             r"""
@@ -33,7 +36,7 @@ def setUp(self) -> None:
         """
         )
 
-        self.preql.add(self.now.shift(seconds=-3), "2 seconds ago")
+        self.preql.add(self.now.shift(seconds=-3), "3 seconds ago")
         self.preql.commit()
 
     def tearDown(self) -> None:
@@ -47,7 +50,20 @@ def tearDown(self) -> None:
     def test_api(self):
         t1 = connect_to_table(TEST_MYSQL_CONN_STRING, "test_api")
         t2 = connect_to_table(TEST_MYSQL_CONN_STRING, ("test_api_2",))
-        assert len(list(diff_tables(t1, t2))) == 1
+        diff = list(diff_tables(t1, t2))
+        assert len(diff) == 1
+
+        t1.database.close()
+        t2.database.close()
+
+        # test where
+        diff_id = diff[0][1][0]
+        where = f"id != {diff_id}"
+
+        t1 = connect_to_table(TEST_MYSQL_CONN_STRING, "test_api", where=where)
+        t2 = connect_to_table(TEST_MYSQL_CONN_STRING, "test_api_2", where=where)
+        diff = list(diff_tables(t1, t2))
+        assert len(diff) == 0
 
         t1.database.close()
         t2.database.close()