dcs-sdk 1.6.0__py3-none-any.whl → 1.6.2__py3-none-any.whl

data_diff/databases/redis.py

@@ -0,0 +1,93 @@
+#  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 os
+import threading
+import time
+
+from loguru import logger
+from redis import ConnectionPool, Redis, RedisError
+
+
+class RedisBackend:
+    _INSTANCE = None
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls, force_recreate: bool = False):
+        with cls._lock:
+            if cls._INSTANCE is None or force_recreate:
+                logger.info("Creating a new RedisBackend instance (force_recreate=%s)", force_recreate)
+                cls._INSTANCE = cls._create_instance()
+            return cls._INSTANCE
+
+    @classmethod
+    def _create_instance(cls):
+        redis_url = (
+            os.getenv("DCS_REDIS_URL", None)
+            or os.getenv("DCS_RABBIT_URL", None)
+            or os.getenv("REDIS_URL", None)
+            or os.getenv("RABBIT_URL", None)
+        )
+        if not redis_url:
+            logger.warning("environment variable is not configured for redis")
+            return None
+        try:
+            pool = ConnectionPool.from_url(
+                redis_url,
+                max_connections=int(os.getenv("DCS_REDIS_MAX_CONNECTIONS", "100")),
+                health_check_interval=int(os.getenv("DCS_REDIS_HEALTH_CHECK_INTERVAL", "30")),
+                socket_connect_timeout=float(os.getenv("DCS_REDIS_SOCKET_CONNECT_TIMEOUT", "5")),
+                socket_timeout=float(os.getenv("DCS_REDIS_SOCKET_TIMEOUT", "5")),
+                socket_keepalive=True,
+            )
+            client = Redis(connection_pool=pool, decode_responses=False, retry_on_timeout=True)
+            # ping to ensure connection and raise early if config wrong
+            client.ping()
+            return cls(client)
+        except RedisError as exc:
+            logger.exception("Failed to create Redis client: %s", exc)
+            raise
+
+    def __init__(self, client: Redis):
+        self._client = client
+
+    @property
+    def client(self) -> Redis:
+        return self._client
+
+    def ensure_connected(self) -> bool:
+        try:
+            self._client.ping()
+            return True
+        except RedisError:
+            logger.warning("Redis ping failed; trying to recreate connection/pool")
+            try:
+                time.sleep(0.2)
+                new_instance = self.__class__._create_instance()
+                self.__class__._INSTANCE = new_instance
+                self._client = new_instance.client
+                self._client.ping()
+                logger.info("Recreated Redis client/pool successfully")
+                return True
+            except Exception as exc:
+                logger.exception("Failed to recreate Redis client: %s", exc)
+                return False
+
+    def close(self):
+        try:
+            self._client.connection_pool.disconnect()
+            logger.info("Closed Redis connection pool")
+        except Exception as e:
+            logger.exception("Error closing Redis connection pool: %s", e)

data_diff/hashdiff_tables.py

@@ -871,7 +871,7 @@ class HashDiffer(HashDiffer):
         level: int,
     ) -> List:
         """Download segment rows and perform in-memory diff."""
-
+        start_time = time.monotonic()
         mode_label = "[IN-MEMORY]" if self.in_memory_diff else "[STANDARD]"
         logger.info(
             ". " * level + f"{mode_label} Downloading rows for comparison: " f"{table1.min_key}..{table1.max_key}"
@@ -887,6 +887,7 @@ class HashDiffer(HashDiffer):
         logger.info(
             ". " * level + f"{mode_label} Downloaded {len(rows1)} and {len(rows2)} rows. "
             f"Total downloaded: {self.stats['rows_downloaded']}"
+            f"Time taken in ms: {int((time.monotonic() - start_time) * 1000)}ms"
         )
 
         # Perform in-memory diff

data_diff/table_segment.py

@@ -25,6 +25,7 @@ from typing_extensions import Self
 
 from data_diff.abcs.database_types import DbKey, DbPath, DbTime, IKey, NumericType
 from data_diff.databases.base import Database
+from data_diff.databases.redis import RedisBackend
 from data_diff.queries.api import (
     SKIP,
     Code,
@@ -50,6 +51,7 @@ from data_diff.utils import (
     ArithTimestamp,
     ArithTimestampTZ,
     ArithUnicodeString,
+    JobCancelledError,
     Vector,
     safezip,
     split_space,
@@ -220,6 +222,7 @@ class TableSegment:
             },
         }
     )
+    job_id: Optional[int] = None
 
     def __attrs_post_init__(self) -> None:
         if not self.update_column and (self.min_update or self.max_update):
@@ -326,7 +329,8 @@ class TableSegment:
 
         # fetched_cols = [NormalizeAsString(this[c]) for c in self.relevant_columns]
         # select = self.make_select().select(*fetched_cols)
-
+        if self._is_cancelled():
+            raise JobCancelledError(self.job_id)
         select = self.make_select().select(*self._relevant_columns_repr)
         start_time = time.monotonic()
         result = self.database.query(select, List[Tuple])
@@ -366,6 +370,8 @@ class TableSegment:
         Returns:
             list: List of tuples containing the queried row data.
         """
+        if self._is_cancelled():
+            raise JobCancelledError(self.job_id)
         select = self.make_select().select(*self._relevant_columns_repr)
 
         filters = []
@@ -471,6 +477,8 @@ class TableSegment:
 
     def count(self) -> int:
         """Count how many rows are in the segment, in one pass."""
+        if self._is_cancelled():
+            raise JobCancelledError(self.job_id)
         start_time = time.monotonic()
         result = self.database.query(self.make_select().select(Count()), int)
         query_time_ms = (time.monotonic() - start_time) * 1000
@@ -480,7 +488,8 @@ class TableSegment:
 
     def count_and_checksum(self) -> Tuple[int, int]:
         """Count and checksum the rows in the segment, in one pass."""
-
+        if self._is_cancelled():
+            raise JobCancelledError(self.job_id)
         checked_columns = [c for c in self.relevant_columns if c not in self.ignored_columns]
         # Build transformed expressions for checksum, honoring transforms and normalization
         checksum_exprs: List[Expr] = []
@@ -552,3 +561,23 @@ class TableSegment:
                 return row_count if row_count is not None else self.count()
         else:
             return row_count if row_count is not None else self.count()
+
+    def _is_cancelled(self) -> bool:
+        run_id = self.job_id
+        if not run_id:
+            return False
+        run_id = f"revoke_job:{run_id}"
+        try:
+            backend = RedisBackend.get_instance()
+            val = backend.client.get(run_id)
+            if not val:
+                return False
+            if isinstance(val, bytes):
+                try:
+                    val = val.decode()
+                except Exception:
+                    val = str(val)
+            return isinstance(val, str) and val.strip().lower() == "revoke"
+        except Exception:
+            logger.warning("Unable to query Redis for cancellation for run_id=%s", run_id)
+            return False

data_diff/utils.py

@@ -1014,3 +1014,9 @@ class SybaseDriverTypes:
     is_ase: bool = False
     is_iq: bool = False
     is_freetds: bool = False
+
+
+class JobCancelledError(RuntimeError):
+    def __init__(self, job_id: str):
+        super().__init__(f"Job ID {job_id} has been revoked.")
+        self.job_id = job_id

dcs_sdk/__version__.py

@@ -12,4 +12,4 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-__version__ = "1.6.0"
+__version__ = "1.6.2"

dcs_sdk/sdk/config/config_loader.py

@@ -86,6 +86,7 @@ class DiffAdvancedConfig(BaseModel):
 
 class Comparison(BaseModel):
     comparison_name: str
+    job_id: Optional[int] = None
     source: SourceTargetConnection
     target: SourceTargetConnection
     source_columns: Optional[List[str]] = None
@@ -429,6 +430,7 @@ class DataDiffConfig:
             )
             new_comparison = {
                 "comparison_name": comparison_name,
+                "job_id": comparison_data.get("job_id", None),
                 "source": self.create_connection_config(
                     source_connection,
                     comparison_data,

dcs_sdk/sdk/data_diff/data_differ.py

@@ -25,6 +25,7 @@ from rich.console import Console
 
 from data_diff import TableSegment, connect, connect_to_table, diff_tables
 from data_diff.databases import Database
+from data_diff.databases.redis import RedisBackend
 from dcs_sdk.sdk.config.config_loader import Comparison, SourceTargetConnection
 from dcs_sdk.sdk.rules.rules_repository import RulesRepository
 from dcs_sdk.sdk.utils.serializer import serialize_table_schema
@@ -160,6 +161,7 @@ class DBTableDiffer:
             extra_columns=tuple(columns),
             where=where,
             transform_columns=config.transform_columns,
+            job_id=self.config.job_id,
         )
 
     def connect_to_db(self, config: SourceTargetConnection, is_source: bool):
@@ -322,7 +324,6 @@ class DBTableDiffer:
 
             self.table1 = self.connect_to_db_table(self.config.source, is_source=True)
             self.table2 = self.connect_to_db_table(self.config.target, is_source=False)
-
             table_1_sample_data = []
             table_2_sample_data = []
             db1_name = (
@@ -399,8 +400,8 @@ class DBTableDiffer:
                 table_1_sample_data = self.table1.with_schema().get_sample_data(limit=100)
                 sample_keys = [list(row[:pks_len]) for row in table_1_sample_data]
                 table_2_sample_data = self.table2.with_schema().get_sample_data(limit=100, sample_keys=sample_keys)
-                if self.config.advanced_configuration.in_memory_diff:
-                    self.config.advanced_configuration.egress_limit = min(max_row_count, 50_00_000)
+                # if self.config.advanced_configuration.in_memory_diff:
+                #     self.config.advanced_configuration.egress_limit = min(max_row_count, 50_00_000)
                 self.diff_iter = diff_tables(
                     self.table1,
                     self.table2,
@@ -658,7 +659,15 @@ class DBTableDiffer:
                                 break
 
                     def sort_key(row, key_columns, extra_columns):
-                        return tuple(row[k] for k in key_columns + extra_columns if k in row)
+                        key_values = []
+                        for k in key_columns + extra_columns:
+                            if k in row:
+                                value = row[k]
+                                if value is None:
+                                    key_values.append("None")
+                                else:
+                                    key_values.append(value)
+                        return tuple(key_values)
 
                     remaining_sources = [row for i, row in enumerate(source_rows) if i not in used_sources]
                     remaining_targets = [row for j, row in enumerate(target_rows) if j not in used_targets]
@@ -745,6 +754,8 @@ class DBTableDiffer:
 
         safe_close(self.source_db)
         safe_close(self.target_db)
+        if self.config.job_id:
+            safe_close(RedisBackend.get_instance())
 
     def cleanup_duckdb(self, src: str, target: str):
         if src and src.endswith("duckdb"):

{dcs_sdk-1.6.0.dist-info → dcs_sdk-1.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dcs-sdk
-Version: 1.6.0
+Version: 1.6.2
 Summary: SDK for DataChecks
 Author: Waterdip Labs
 Author-email: hello@waterdip.ai
@@ -49,6 +49,7 @@ Requires-Dist: pydantic (>=1.10.12)
 Requires-Dist: pyodbc (>=4.0.39) ; extra == "mssql" or extra == "sybase" or extra == "all-dbs"
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
+Requires-Dist: redis[hiredis] (>=5.2.1,<6.0.0)
 Requires-Dist: requests (>=2.32.4,<3.0.0)
 Requires-Dist: rich (>=13.8.0)
 Requires-Dist: snowflake-connector-python (>=3.17.2) ; extra == "snowflake" or extra == "all-dbs"
@@ -61,7 +62,7 @@ Requires-Dist: vertica-python (>=1.4.0) ; extra == "vertica" or extra == "all-db
 Description-Content-Type: text/markdown
 
 <h1 align="center">
-  DCS SDK v1.6.0
+  DCS SDK v1.6.2
 </h1>
 
 > SDK for DataChecks

{dcs_sdk-1.6.0.dist-info → dcs_sdk-1.6.2.dist-info}/RECORD

@@ -16,6 +16,7 @@ data_diff/databases/mysql.py,sha256=LTuLTkxac16y7w5_q1nEuKy6ZefqQUBQqV_tdnzjNT4,
 data_diff/databases/oracle.py,sha256=GIczARg96lOlmgUpHd-hOFl8etFJmxWMy-KkA2COG74,8031
 data_diff/databases/postgresql.py,sha256=dSZBUYbGwWnvCwTxSwMSiaElmQSpB2Q3chsmgvUeqQE,10787
 data_diff/databases/presto.py,sha256=1Z8iDV2pc35Bu7DuUerFuFLbrwgSHSkBYmJ72JlZSZ8,7116
+data_diff/databases/redis.py,sha256=gspPgUr7uFyZrqhIHBgYZo7-ecKNn8Wc_S8CVSQasLU,3501
 data_diff/databases/redshift.py,sha256=-gFWs3NCcevO4s6c4zV3_LYihK24fUd5BADTKahubjw,8122
 data_diff/databases/snowflake.py,sha256=7G6fvVJXOtTvXmSfWCxTslF4WohoscQoiqcmJIN684A,7910
 data_diff/databases/sybase.py,sha256=RCZu2sqY08w8CN5NY0WN6QZFJBvQ3Ezk8OtfH8enxaY,31311
@@ -24,7 +25,7 @@ data_diff/databases/vertica.py,sha256=2dSDZp6qOEvUVPldI5Tgn7Sm3dCpC3vNXJL3qb3FDv
 data_diff/diff_tables.py,sha256=AyFJF6oaam06AH4ZPI8pj63BiYojHoZTykjrdJCX2fI,20899
 data_diff/errors.py,sha256=4Yru8yOwyuDuBlTABnGCvJMSpe6-rbLJpNnVHeTTyHU,745
 data_diff/format.py,sha256=QFDjdZaBVf_N-jfKiX4ppOUdpXTPZXmv1j0pc1RiOoc,10245
-data_diff/hashdiff_tables.py,sha256=oElu0rFAFaRGVZpiIl546yUVYQm02J6HFjVkpuHHFIQ,43134
+data_diff/hashdiff_tables.py,sha256=c-YOOkjUGUok_MJpvkwVOxfJPIQdvtJkW686kPLUYsM,43252
 data_diff/info_tree.py,sha256=yHtFSoXuu6oBafLYOYQjUSKlB-DnAAd08U9HOEAdTPI,2799
 data_diff/joindiff_tables.py,sha256=fyrEYjyh2BX1vGibwVZLYM1V6JJTOY-uGXY-KInvMkw,17612
 data_diff/lexicographic_space.py,sha256=bBoCbbH1Mla9jNOq1b5RuwjAxSVU7gWkra673tPBwXQ,8305
@@ -37,20 +38,20 @@ data_diff/queries/base.py,sha256=pT-iaII7Nlu-w-Cuq9fhoNKX7-GSxkQ3Fk8K-tMkk60,964
 data_diff/queries/extras.py,sha256=aUm-ifj3BMlz4o4bbuHtmnvHZuptYAKGS5yWTHmNpvc,1270
 data_diff/query_utils.py,sha256=R7ZfRwcvv9Zf4zWXNln4tr_OxLmDI7CPmmCahYfHxlo,2101
 data_diff/schema.py,sha256=QoYSSB3k-svLXz680uRgsI4qjii8BFKOOQvheqtgEbs,2413
-data_diff/table_segment.py,sha256=SdBVHLndDWMuLLM9zuUbfyTXZtNQgXarVLOmQTUWCCY,23356
+data_diff/table_segment.py,sha256=pGSrTmtVt4ioDj06U47XBlJf_p5EtwO2QVXR9iunkZM,24491
 data_diff/thread_utils.py,sha256=_692ERjnWfHKaZsLdg7CNfkKiRd66y7_kpgDwzntp44,3831
-data_diff/utils.py,sha256=ZSgK6utdYDWvnsYgl7sJQa6J3XK3T6xSB_BGdQVeK24,33457
+data_diff/utils.py,sha256=A0ZoxIngdU-uiiFp1uXcNx1XlEsLWgXr8J36my_ViSw,33627
 data_diff/version.py,sha256=Wk0ovyBlLEF2UaWLWEcVBLFElREtIxi7TU1hD3CuTFI,634
 dcs_sdk/__init__.py,sha256=RkfhRKLXEForLCs4rZkTf0qc_b0TokSggSAcKI4yfZg,610
 dcs_sdk/__main__.py,sha256=Qn8stIaQGrdLjHQ-H7xO0T-brtq5RWZoWU9QvqoarV8,683
-dcs_sdk/__version__.py,sha256=Ou4IXxHmCirdCoZjm7td7gC94Dl0xUFPJxDuZZwlVvQ,633
+dcs_sdk/__version__.py,sha256=Pa1cOhyUyQyETiOwXfxJyAF2FKshzqr1qkVSsgMjYAs,633
 dcs_sdk/cli/__init__.py,sha256=RkfhRKLXEForLCs4rZkTf0qc_b0TokSggSAcKI4yfZg,610
 dcs_sdk/cli/cli.py,sha256=LyrRk972OL9pTqrvBeXWBu5rUDAN17lQ1g8FdSRW_8M,4299
 dcs_sdk/sdk/__init__.py,sha256=skrZcgWWJBL6NXTUERywJ3qRJRemgpDXyW7lPg1FJk8,2107
 dcs_sdk/sdk/config/__init__.py,sha256=RkfhRKLXEForLCs4rZkTf0qc_b0TokSggSAcKI4yfZg,610
-dcs_sdk/sdk/config/config_loader.py,sha256=hVdGidcm6ibqJ0yt7fZcL_0kFcorBfM7vHGhUiBH9Ww,21587
+dcs_sdk/sdk/config/config_loader.py,sha256=oooSTV6QjbXKpCkwpl6vcBdjABGT-h99vBbWTbIkmjc,21683
 dcs_sdk/sdk/data_diff/__init__.py,sha256=RkfhRKLXEForLCs4rZkTf0qc_b0TokSggSAcKI4yfZg,610
-dcs_sdk/sdk/data_diff/data_differ.py,sha256=d8hwXs0zG5AuaRjsNlpFShpTxA1_YPgH2aCRsJQfDMU,35846
+dcs_sdk/sdk/data_diff/data_differ.py,sha256=zxWa-mYAdfZepNuXz1h_xxFQBC4tdhBqlbZCVEfb8Y8,36378
 dcs_sdk/sdk/rules/__init__.py,sha256=_BkKcE_jfdDQI_ECdOamJaefMKEXrKpYjPpnBQXl_Xs,657
 dcs_sdk/sdk/rules/rules_mappping.py,sha256=fxakVkf7B2cVkYSO946LTim_HmMsl6lBDBqZjTTsSPI,1292
 dcs_sdk/sdk/rules/rules_repository.py,sha256=x0Rli-wdnHAmXm5526go_qC3P-eFRt-4L7fs4hNqC-g,7564
@@ -65,7 +66,7 @@ dcs_sdk/sdk/utils/similarity_score/levenshtein_distance_provider.py,sha256=puAWP
 dcs_sdk/sdk/utils/table.py,sha256=X8HxdYTWyx_oVrBWPsXlmA-xJKXXDBW9RrhlWNqA1As,18224
 dcs_sdk/sdk/utils/themes.py,sha256=Meo2Yldv4uyPpEqI7qdA28Aa6sxtwUU1dLKKm4QavjM,1403
 dcs_sdk/sdk/utils/utils.py,sha256=vF2zAvgt__Y8limicWTEWRyn41SBVJN81ZCTBRy6hQg,11907
-dcs_sdk-1.6.0.dist-info/METADATA,sha256=lYbmpwlrScvWP06pR_YYH511rTcqoF5A6GJurZ3jubc,6275
-dcs_sdk-1.6.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-dcs_sdk-1.6.0.dist-info/entry_points.txt,sha256=zQtrZL7YuaKtt6WPwihCTV1BRXnqBkaY6zUGdYJbBSg,49
-dcs_sdk-1.6.0.dist-info/RECORD,,
+dcs_sdk-1.6.2.dist-info/METADATA,sha256=pSN0LCAuwyoWjSaHN_pibtgtHKLH8nsea-A2sbgdeRg,6322
+dcs_sdk-1.6.2.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+dcs_sdk-1.6.2.dist-info/entry_points.txt,sha256=zQtrZL7YuaKtt6WPwihCTV1BRXnqBkaY6zUGdYJbBSg,49
+dcs_sdk-1.6.2.dist-info/RECORD,,