Joindiff: Fix stats collections

Author: erezsh

data_diff/databases/database_types.py

@@ -141,6 +141,8 @@ class UnknownColType(ColType):
 
 
 class AbstractDialect(ABC):
+ """Dialect-dependent query expressions"""
+
  name: str
 
  @abstractmethod
@@ -177,56 +179,18 @@ def explain_as_text(self, query: str) -> str:
  "Provide SQL for explaining a query, returned in as table(varchar)"
  ...
 
-
-class AbstractDatabase(AbstractDialect):
  @abstractmethod
- def timestamp_value(self, t: DbTime) -> str:
+ def timestamp_value(self, t: datetime) -> str:
  "Provide SQL for the given timestamp value"
  ...
 
- @abstractmethod
- def md5_to_int(self, s: str) -> str:
- "Provide SQL for computing md5 and returning an int"
- ...
-
- @abstractmethod
- def _query(self, sql_code: str) -> list:
- "Send query to database and return result"
- ...
-
- @abstractmethod
- def select_table_schema(self, path: DbPath) -> str:
- "Provide SQL for selecting the table schema as (name, type, date_prec, num_prec)"
- ...
 
- @abstractmethod
- def query_table_schema(self, path: DbPath) -> Dict[str, tuple]:
- """Query the table for its schema for table in 'path', and return {column: tuple}
- where the tuple is (table_name, col_name, type_repr, datetime_precision?, numeric_precision?, numeric_scale?)
- """
- ...
+class AbstractDatadiffDialect(ABC):
+ """Dialect-dependent query expressions, that are specific to data-diff"""
 
  @abstractmethod
- def _process_table_schema(
- self, path: DbPath, raw_schema: Dict[str, tuple], filter_columns: Sequence[str], where: str = None
- ):
- """Process the result of query_table_schema().
-
- Done in a separate step, to minimize the amount of processed columns.
- Needed because processing each column may:
- * throw errors and warnings
- * query the database to sample values
-
- """
-
- @abstractmethod
- def parse_table_name(self, name: str) -> DbPath:
- "Parse the given table name into a DbPath"
- ...
-
- @abstractmethod
- def close(self):
- "Close connection(s) to the database instance. Querying will stop functioning."
+ def md5_to_int(self, s: str) -> str:
+ "Provide SQL for computing md5 and returning an int"
  ...
 
  @abstractmethod
@@ -294,6 +258,48 @@ def normalize_value_by_type(self, value: str, coltype: ColType) -> str:
  return self.normalize_uuid(value, coltype)
  return self.to_string(value)
 
+
+class AbstractDatabase(AbstractDialect, AbstractDatadiffDialect):
+ @abstractmethod
+ def _query(self, sql_code: str) -> list:
+ "Send query to database and return result"
+ ...
+
+ @abstractmethod
+ def select_table_schema(self, path: DbPath) -> str:
+ "Provide SQL for selecting the table schema as (name, type, date_prec, num_prec)"
+ ...
+
+ @abstractmethod
+ def query_table_schema(self, path: DbPath) -> Dict[str, tuple]:
+ """Query the table for its schema for table in 'path', and return {column: tuple}
+ where the tuple is (table_name, col_name, type_repr, datetime_precision?, numeric_precision?, numeric_scale?)
+ """
+ ...
+
+ @abstractmethod
+ def _process_table_schema(
+ self, path: DbPath, raw_schema: Dict[str, tuple], filter_columns: Sequence[str], where: str = None
+ ):
+ """Process the result of query_table_schema().
+
+ Done in a separate step, to minimize the amount of processed columns.
+ Needed because processing each column may:
+ * throw errors and warnings
+ * query the database to sample values
+
+ """
+
+ @abstractmethod
+ def parse_table_name(self, name: str) -> DbPath:
+ "Parse the given table name into a DbPath"
+ ...
+
+ @abstractmethod
+ def close(self):
+ "Close connection(s) to the database instance. Querying will stop functioning."
+ ...
+
  @abstractmethod
  def _normalize_table_path(self, path: DbPath) -> DbPath:
  ...

data_diff/databases/presto.py

@@ -83,7 +83,7 @@ def close(self):
  self._conn.close()
 
  def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
- # TODO
+ # TODO rounds
  if coltype.rounds:
  s = f"date_format(cast({value} as timestamp(6)), '%Y-%m-%d %H:%i:%S.%f')"
  else:

data_diff/joindiff_tables.py

@@ -10,7 +10,7 @@
 
 from runtype import dataclass
 
-from data_diff.databases.database_types import DbPath, Schema
+from data_diff.databases.database_types import DbPath, NumericType, Schema
 from data_diff.databases.base import QueryError
 
 
@@ -273,7 +273,7 @@ def _collect_stats(self, i, table):
  f"max_{c}": max_(this[c]),
  }
  for c in table._relevant_columns
- if c == "id" # TODO just if the right type
+ if isinstance(table._schema[c], NumericType)
  )
  col_exprs["count"] = Count()
 

data_diff/table_segment.py

@@ -177,7 +177,6 @@ def count_and_checksum(self) -> Tuple[int, int]:
  def query_key_range(self) -> Tuple[int, int]:
  """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
- # TODO better error if there is no schema
  (k,) = self.key_columns
  select = self._make_select().select(
  ApplyFuncAndNormalizeAsString(this[k], min_),

tests/test_diff_tables.py

@@ -485,8 +485,6 @@ def setUp(self):
  self.new_uuid = uuid.uuid1(32132131)
  queries.append(f"INSERT INTO {self.table_src} VALUES ('{self.new_uuid}', 'This one is different')")
 
- # TODO test unexpected values?
-
  for query in queries:
  self.connection.query(query, None)
 
@@ -542,8 +540,6 @@ def setUp(self):
  self.new_alphanum = "aBcDeFgHiJ"
  queries.append(f"INSERT INTO {self.table_src} VALUES ('{self.new_alphanum}', 'This one is different')")
 
- # TODO test unexpected values?
-
  for query in queries:
  self.connection.query(query, None)
 
@@ -594,8 +590,6 @@ def setUp(self):
  self.new_alphanum = "aBcDeFgHiJ"
  queries.append(f"INSERT INTO {self.table_src} VALUES ('{self.new_alphanum}', 'This one is different')")
 
- # TODO test unexpected values?
-
  for query in queries:
  self.connection.query(query, None)
 

tests/test_joindiff.py

@@ -137,6 +137,8 @@ def test_diff_small_tables(self):
  self.assertEqual(expected, diff)
  self.assertEqual(2, self.differ.stats["table1_count"])
  self.assertEqual(1, self.differ.stats["table2_count"])
+ self.assertEqual(3, self.differ.stats["table1_sum_id"])
+ self.assertEqual(1, self.differ.stats["table2_sum_id"])
 
  # Test materialize
  materialize_path = self.connection.parse_table_name(f"test_mat_{random_table_suffix()}")

tests/test_query.py

@@ -1,4 +1,4 @@
-from cmath import exp
+from datetime import datetime
 from typing import List, Optional
 import unittest
 from data_diff.databases.database_types import AbstractDialect, CaseInsensitiveDict, CaseSensitiveDict
@@ -35,6 +35,9 @@ def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None
  def explain_as_text(self, query: str) -> str:
  return f"explain {query}"
 
+ def timestamp_value(self, t: datetime) -> str:
+ return f"timestamp '{t}'"
+
 
 class TestQuery(unittest.TestCase):
  def setUp(self):