Cleanup and minor fixes (pylint pass)

Author: erezsh

data_diff/__init__.py

@@ -70,24 +70,27 @@ def diff_tables(
 
     Parameters:
         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`.
+        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:`DbKey`, optional): Lowest key value, used to restrict the segment
         max_key (:data:`DbKey`, 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
         algorithm (:class:`Algorithm`): Which diffing algorithm to use (`HASHDIFF` or `JOINDIFF`)
-        bisection_factor (int): Into how many segments to bisect per iteration. (when algorithm is `HASHDIFF`)
-        bisection_threshold (Number): When should we stop bisecting and compare locally (when algorithm is `HASHDIFF`; in row count).
+        bisection_factor (int): Into how many segments to bisect per iteration. (Used when algorithm is `HASHDIFF`)
+        bisection_threshold (Number): Minimal row count of segment to bisect, otherwise download
+                                      and compare locally. (Used when algorithm is `HASHDIFF`).
         threaded (bool): Enable/disable threaded diffing. Needed to take advantage of database threads.
-        max_threadpool_size (int): Maximum size of each threadpool. ``None`` means auto. Only relevant when `threaded` is ``True``.
+        max_threadpool_size (int): Maximum size of each threadpool. ``None`` means auto.
+                                   Only relevant when `threaded` is ``True``.
                                    There may be many pools, so number of actual threads can be a lot higher.
 
     Note:
         The following parameters are used to override the corresponding attributes of the given :class:`TableSegment` instances:
-        `key_columns`, `update_column`, `extra_columns`, `min_key`, `max_key`. If different values are needed per table, it's
-        possible to omit them here, and instead set them directly when creating each :class:`TableSegment`.
+        `key_columns`, `update_column`, `extra_columns`, `min_key`, `max_key`.
+        If different values are needed per table, it's possible to omit them here, and instead set
+        them directly when creating each :class:`TableSegment`.
 
     Example:
         >>> table1 = connect_to_table('postgresql:///', 'Rating', 'id')

data_diff/databases/base.py

@@ -8,6 +8,7 @@
 from abc import abstractmethod
 
 from data_diff.utils import is_uuid, safezip
+from data_diff.queries import Expr, Compiler, table, Select, SKIP, Explain
 from .database_types import (
     AbstractDatabase,
     ColType,
@@ -27,8 +28,6 @@
     DbPath,
 )
 
-from data_diff.queries import Expr, Compiler, table, Select, SKIP, Explain
-
 logger = logging.getLogger("database")
 
 
@@ -110,6 +109,8 @@ class Database(AbstractDatabase):
     default_schema: str = None
     SUPPORTS_ALPHANUMS = True
 
+    _interactive = False
+
     @property
     def name(self):
         return type(self).__name__
@@ -126,11 +127,14 @@ def query(self, sql_ast: Union[Expr, Generator], res_type: type = list):
                 return SKIP
 
         logger.debug("Running SQL (%s): %s", type(self).__name__, sql_code)
-        if getattr(self, "_interactive", False) and isinstance(sql_ast, Select):
+        if self._interactive and isinstance(sql_ast, Select):
             explained_sql = compiler.compile(Explain(sql_ast))
             explain = self._query(explained_sql)
-            for (row,) in explain:
-                logger.debug(f"EXPLAIN: {row}")
+            for row in explain:
+                # Most returned a 1-tuple. Presto returns a string
+                if isinstance(row, tuple):
+                    row ,= row
+                logger.debug("EXPLAIN: %s", row)
             answer = input("Continue? [y/n] ")
             if not answer.lower() in ["y", "yes"]:
                 sys.exit(1)

data_diff/databases/bigquery.py

@@ -1,4 +1,5 @@
-from .database_types import *
+from typing import Union
+from .database_types import Timestamp, Datetime, Integer, Decimal, Float, Text, DbPath, FractionalType, TemporalType
 from .base import Database, import_helper, parse_table_name, ConnectError, apply_query
 from .base import TIMESTAMP_PRECISION_POS, ThreadLocalInterpreter
 

data_diff/databases/database_types.py

@@ -1,6 +1,6 @@
 import logging
 import decimal
-from abc import ABC, abstractmethod, abstractproperty
+from abc import ABC, abstractmethod
 from typing import Sequence, Optional, Tuple, Union, Dict, List
 from datetime import datetime
 
@@ -234,30 +234,6 @@ def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
         """
         ...
 
-    def normalize_value_by_type(self, value: str, coltype: ColType) -> str:
-        """Creates an SQL expression, that converts 'value' to a normalized representation.
-
-        The returned expression must accept any SQL value, and return a string.
-
-        The default implementation dispatches to a method according to `coltype`:
-
-        ::
-
-            TemporalType    -> normalize_timestamp()
-            FractionalType  -> normalize_number()
-            *else*          -> to_string()
-
-            (`Integer` falls in the *else* category)
-
-        """
-        if isinstance(coltype, TemporalType):
-            return self.normalize_timestamp(value, coltype)
-        elif isinstance(coltype, FractionalType):
-            return self.normalize_number(value, coltype)
-        elif isinstance(coltype, ColType_UUID):
-            return self.normalize_uuid(value, coltype)
-        return self.to_string(value)
-
 
 class AbstractDatabase(AbstractDialect, AbstractDatadiffDialect):
     @abstractmethod
@@ -304,10 +280,35 @@ def close(self):
     def _normalize_table_path(self, path: DbPath) -> DbPath:
         ...
 
-    @abstractproperty
+    @property
+    @abstractmethod
     def is_autocommit(self) -> bool:
         ...
 
+    def normalize_value_by_type(self, value: str, coltype: ColType) -> str:
+        """Creates an SQL expression, that converts 'value' to a normalized representation.
+
+        The returned expression must accept any SQL value, and return a string.
+
+        The default implementation dispatches to a method according to `coltype`:
+
+        ::
+
+            TemporalType    -> normalize_timestamp()
+            FractionalType  -> normalize_number()
+            *else*          -> to_string()
+
+            (`Integer` falls in the *else* category)
+
+        """
+        if isinstance(coltype, TemporalType):
+            return self.normalize_timestamp(value, coltype)
+        elif isinstance(coltype, FractionalType):
+            return self.normalize_number(value, coltype)
+        elif isinstance(coltype, ColType_UUID):
+            return self.normalize_uuid(value, coltype)
+        return self.to_string(value)
+
 
 Schema = CaseAwareMapping
 

data_diff/databases/databricks.py

@@ -1,6 +1,18 @@
+from typing import Dict, Sequence
 import logging
 
-from .database_types import *
+from .database_types import (
+    Integer,
+    Float,
+    Decimal,
+    Timestamp,
+    Text,
+    TemporalType,
+    NumericType,
+    DbPath,
+    ColType,
+    UnknownColType,
+)
 from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, Database, import_helper, parse_table_name
 
 

data_diff/databases/mysql.py

@@ -1,4 +1,14 @@
-from .database_types import *
+from .database_types import (
+    Datetime,
+    Timestamp,
+    Float,
+    Decimal,
+    Integer,
+    Text,
+    TemporalType,
+    FractionalType,
+    ColType_UUID,
+)
 from .base import ThreadedDatabase, import_helper, ConnectError
 from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, TIMESTAMP_PRECISION_POS
 

data_diff/databases/oracle.py

@@ -1,6 +1,20 @@
+from typing import Dict, List, Optional
+
 from ..utils import match_regexps
 
-from .database_types import *
+from .database_types import (
+    Decimal,
+    Float,
+    Text,
+    DbPath,
+    TemporalType,
+    ColType,
+    DbTime,
+    ColType_UUID,
+    Timestamp,
+    TimestampTZ,
+    FractionalType,
+)
 from .base import ThreadedDatabase, import_helper, ConnectError, QueryError
 from .base import TIMESTAMP_PRECISION_POS
 

data_diff/databases/postgresql.py

@@ -1,4 +1,14 @@
-from .database_types import *
+from .database_types import (
+    Timestamp,
+    TimestampTZ,
+    Float,
+    Decimal,
+    Integer,
+    TemporalType,
+    Native_UUID,
+    Text,
+    FractionalType,
+)
 from .base import ThreadedDatabase, import_helper, ConnectError
 from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, _CHECKSUM_BITSIZE, TIMESTAMP_PRECISION_POS
 

data_diff/databases/presto.py

@@ -3,7 +3,19 @@
 
 from data_diff.utils import match_regexps
 
-from .database_types import *
+from .database_types import (
+    Timestamp,
+    TimestampTZ,
+    Integer,
+    Float,
+    Text,
+    FractionalType,
+    DbPath,
+    Decimal,
+    ColType,
+    ColType_UUID,
+    TemporalType,
+)
 from .base import Database, import_helper, ThreadLocalInterpreter
 from .base import (
     MD5_HEXDIGITS,
@@ -17,7 +29,7 @@ def query_cursor(c, sql_code):
     if sql_code.lower().startswith("select"):
         return c.fetchall()
     # Required for the query to actually run 🤯
-    if re.match(r"(insert|create|truncate|drop)", sql_code, re.IGNORECASE):
+    if re.match(r"(insert|create|truncate|drop|explain)", sql_code, re.IGNORECASE):
         return c.fetchone()
 
 
@@ -98,7 +110,7 @@ def select_table_schema(self, path: DbPath) -> str:
         schema, table = self._normalize_table_path(path)
 
         return (
-            "SELECT column_name, data_type, 3 as datetime_precision, 3 as numeric_precision "
+            "SELECT column_name, data_type, 3 as datetime_precision, 3 as numeric_precision, NULL as numeric_scale "
             "FROM INFORMATION_SCHEMA.COLUMNS "
             f"WHERE table_name = '{table}' AND table_schema = '{schema}'"
         )
@@ -110,6 +122,7 @@ def _parse_type(
         type_repr: str,
         datetime_precision: int = None,
         numeric_precision: int = None,
+        numeric_scale: int = None,
     ) -> ColType:
         timestamp_regexps = {
             r"timestamp\((\d)\)": Timestamp,
@@ -134,5 +147,9 @@ def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
         # Trim doesn't work on CHAR type
         return f"TRIM(CAST({value} AS VARCHAR))"
 
+    @property
     def is_autocommit(self) -> bool:
         return False
+
+    def explain_as_text(self, query: str) -> str:
+        return f"EXPLAIN (FORMAT TEXT) {query}"

data_diff/databases/redshift.py

@@ -1,4 +1,5 @@
-from .database_types import *
+from typing import List
+from .database_types import Float, TemporalType, FractionalType, DbPath
 from .postgresql import PostgreSQL, MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, TIMESTAMP_PRECISION_POS
 
 

data_diff/databases/snowflake.py

@@ -1,6 +1,7 @@
+from typing import Union
 import logging
 
-from .database_types import *
+from .database_types import Timestamp, TimestampTZ, Decimal, Float, Text, FractionalType, TemporalType, DbPath
 from .base import ConnectError, Database, import_helper, CHECKSUM_MASK, ThreadLocalInterpreter
 
 
@@ -88,6 +89,7 @@ def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
     def normalize_number(self, value: str, coltype: FractionalType) -> str:
         return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))")
 
+    @property
     def is_autocommit(self) -> bool:
         return True
 

data_diff/databases/trino.py

@@ -1,4 +1,4 @@
-from .database_types import *
+from .database_types import TemporalType, ColType_UUID
 from .presto import Presto
 from .base import import_helper
 from .base import TIMESTAMP_PRECISION_POS

data_diff/diff_tables.py

@@ -78,6 +78,7 @@ def _run_in_background(self, *funcs):
 
 class TableDiffer(ThreadBase, ABC):
     bisection_factor = 32
+    stats: dict = {}
 
     def diff_tables(self, table1: TableSegment, table2: TableSegment) -> DiffResult:
         """Diff the given tables.
@@ -177,7 +178,6 @@ def _bisect_and_diff_tables(self, table1, table2):
         table1, table2 = [t.new(min_key=min_key1, max_key=max_key1) for t in (table1, table2)]
 
         logger.info(
-            # f"Diffing tables | segments: {self.bisection_factor}, bisection threshold: {self.bisection_threshold}. "
             f"Diffing segments at key-range: {table1.min_key}..{table2.max_key}. "
             f"size: table1 <= {table1.approximate_size()}, table2 <= {table2.approximate_size()}"
         )