Reformatted with 'Black', as requested

Author: erezsh

data-diff/__init__.py

@@ -3,6 +3,7 @@
 from .database import connect_to_uri
 from .diff_tables import TableSegment, TableDiffer
 
+
 def create_source(db_uri: str, table_name: str, key_column: str, extra_columns: Tuple[str, ...] = ()):
  db = connect_to_uri(db_uri)
- return TableSegment(db, (table_name,), key_column, tuple(extra_columns))
+ return TableSegment(db, (table_name,), key_column, tuple(extra_columns))

data-diff/__main__.py

@@ -8,23 +8,37 @@
 
 import click
 
-LOG_FORMAT = '[%(asctime)s] %(levelname)s - %(message)s'
-DATE_FORMAT ='%H:%M:%S'
+LOG_FORMAT = "[%(asctime)s] %(levelname)s - %(message)s"
+DATE_FORMAT = "%H:%M:%S"
+
 
 @click.command()
-@click.argument('db1_uri')
-@click.argument('table1_name')
-@click.argument('db2_uri')
-@click.argument('table2_name')
-@click.option('-k', '--key_column', default='id', help='Name of primary key column')
-@click.option('-c', '--columns', default=['updated_at'], multiple=True, help='Names of extra columns to compare')
-@click.option('-l', '--limit', default=None, help='Maximum number of differences to find')
-@click.option('--bisection-factor', default=32, help='Segments per iteration')
-@click.option('--bisection-threshold', default=1024**2, help='Minimal bisection threshold')
-@click.option('-s', '--stats', is_flag=True, help='Print stats instead of a detailed diff')
-@click.option('-d', '--debug', is_flag=True, help='Print debug info')
-@click.option('-v', '--verbose', is_flag=True, help='Print extra info')
-def main(db1_uri, table1_name, db2_uri, table2_name, key_column, columns, limit, bisection_factor, bisection_threshold, stats, debug, verbose):
+@click.argument("db1_uri")
+@click.argument("table1_name")
+@click.argument("db2_uri")
+@click.argument("table2_name")
+@click.option("-k", "--key_column", default="id", help="Name of primary key column")
+@click.option("-c", "--columns", default=["updated_at"], multiple=True, help="Names of extra columns to compare")
+@click.option("-l", "--limit", default=None, help="Maximum number of differences to find")
+@click.option("--bisection-factor", default=32, help="Segments per iteration")
+@click.option("--bisection-threshold", default=1024**2, help="Minimal bisection threshold")
+@click.option("-s", "--stats", is_flag=True, help="Print stats instead of a detailed diff")
+@click.option("-d", "--debug", is_flag=True, help="Print debug info")
+@click.option("-v", "--verbose", is_flag=True, help="Print extra info")
+def main(
+ db1_uri,
+ table1_name,
+ db2_uri,
+ table2_name,
+ key_column,
+ columns,
+ limit,
+ bisection_factor,
+ bisection_threshold,
+ stats,
+ debug,
+ verbose,
+):
  if limit and stats:
  print("Error: cannot specify a limit when using the -s/--stats switch")
  return
@@ -34,7 +48,6 @@ def main(db1_uri, table1_name, db2_uri, table2_name, key_column, columns, limit,
  elif verbose:
  logging.basicConfig(level=logging.INFO, format=LOG_FORMAT, datefmt=DATE_FORMAT)
 
-
  db1 = connect_to_uri(db1_uri)
  db2 = connect_to_uri(db2_uri)
 
@@ -54,8 +67,8 @@ def main(db1_uri, table1_name, db2_uri, table2_name, key_column, columns, limit,
  percent = 100 * len(diff) / table1.count
  print(f"Diff-Total: {len(diff)} changed rows out of {table1.count}")
  print(f"Diff-Percent: {percent:.4f}%")
- plus = len([1 for op,_ in diff if op=='+'])
- minus = len([1 for op,_ in diff if op=='-'])
+ plus = len([1 for op, _ in diff if op == "+"])
+ minus = len([1 for op, _ in diff if op == "-"])
  print(f"Diff-Split: +{plus} -{minus}")
  else:
  for op, key in diff_iter:
@@ -67,5 +80,5 @@ def main(db1_uri, table1_name, db2_uri, table2_name, key_column, columns, limit,
  logging.info(f"Duration: {end-start:.2f} seconds.")
 
 
-if __name__ == '__main__':
- main()
+if __name__ == "__main__":
+ main()

data-diff/database.py

@@ -7,29 +7,38 @@
 from .sql import SqlOrStr, Compiler
 
 
-logger = logging.getLogger('database')
+logger = logging.getLogger("database")
 
 
 def import_postgres():
  import psycopg2
  import psycopg2.extras
+
  psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
  return psycopg2
 
+
 def import_mysql():
  import mysql.connector
+
  return mysql.connector
 
+
 def import_snowflake():
  import snowflake.connector
+
  return snowflake
 
+
 def import_mssql():
  import pymssql
+
  return pymssql
 
+
 def import_oracle():
  import cx_Oracle
+
  return cx_Oracle
 
 
@@ -38,12 +47,13 @@ class ConnectError(Exception):
 
 
 def _one(seq):
- x ,= seq
+ (x,) = seq
  return x
 
+
 class Database(ABC):
  """Base abstract class for databases.
- 
+
  Used for providing connection code and implementation specific SQL utilities.
  """
 
@@ -60,10 +70,10 @@ def query(self, sql_ast: SqlOrStr, res_type: type):
  res = self._query(sql_code)
  if res_type is int:
  res = _one(_one(res))
- if res is None:  # May happen due to sum() of 0 items
+ if res is None: # May happen due to sum() of 0 items
  return None
  return int(res)
- elif getattr(res_type, '__origin__', None) is list and len(res_type.__args__) == 1:
+ elif getattr(res_type, "__origin__", None) is list and len(res_type.__args__) == 1:
  if res_type.__args__ == (int,):
  return [_one(row) for row in res]
  elif res_type.__args__ == (Tuple,):
@@ -87,12 +97,14 @@ def md5_to_int(self, s: str) -> str:
  "Provide SQL for computing md5 and returning an int"
  ...
 
-CHECKSUM_HEXDIGITS = 15 # Must be 15 or lower
+
+CHECKSUM_HEXDIGITS = 15 # Must be 15 or lower
 MD5_HEXDIGITS = 32
 
-_CHECKSUM_BITSIZE = CHECKSUM_HEXDIGITS<<2
+_CHECKSUM_BITSIZE = CHECKSUM_HEXDIGITS << 2
 CHECKSUM_MASK = (2**_CHECKSUM_BITSIZE) - 1
 
+
 class Postgres(Database):
  def __init__(self, host, port, database, user, password):
  postgres = import_postgres()
@@ -110,18 +122,18 @@ def md5_to_int(self, s: str) -> str:
  return f"('x' || substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}))::bit({_CHECKSUM_BITSIZE})::bigint"
 
  def to_string(self, s: str):
- return f'{s}::varchar'
+ return f"{s}::varchar"
 
 
 class MySQL(Database):
  def __init__(self, host, port, database, user, password):
  mysql = import_mysql()
 
  args = dict(host=host, port=port, database=database, user=user, password=password)
- self._args = {k:v for k, v in args.items() if v is not None}
+ self._args = {k: v for k, v in args.items() if v is not None}
 
  try:
- self._conn = mysql.connect(charset='utf8', use_unicode=True, **self._args)
+ self._conn = mysql.connect(charset="utf8", use_unicode=True, **self._args)
  except mysql.Error as e:
  if e.errno == mysql.errorcode.ER_ACCESS_DENIED_ERROR:
  raise ConnectError("Bad user name or password") from e
@@ -131,13 +143,13 @@ def __init__(self, host, port, database, user, password):
  raise ConnectError(*e.args) from e
 
  def quote(self, s: str):
- return f'`{s}`'
+ return f"`{s}`"
 
  def md5_to_int(self, s: str) -> str:
  return f"cast(conv(substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16, 10) as unsigned)"
 
  def to_string(self, s: str):
- return f'cast({s} as char)' 
+ return f"cast({s} as char)"
 
 
 class Oracle(Database):
@@ -155,31 +167,33 @@ def md5_to_int(self, s: str) -> str:
  return f"to_number(substr(standard_hash({s}, 'MD5'), 18), 'xxxxxxxxxxxxxxx')"
 
  def quote(self, s: str):
- return f'{s}'
+ return f"{s}"
 
  def to_string(self, s: str):
- return f'cast({s} as varchar(1024))' 
+ return f"cast({s} as varchar(1024))"
+
 
 class Redshift(Postgres):
  def md5_to_int(self, s: str) -> str:
  return f"strtol(substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16)::decimal(38)"
 
+
 class MsSQL(Database):
  "AKA sql-server"
 
  def __init__(self, host, port, database, user, password):
  mssql = import_mssql()
 
  args = dict(server=host, port=port, database=database, user=user, password=password)
- self._args = {k:v for k, v in args.items() if v is not None}
+ self._args = {k: v for k, v in args.items() if v is not None}
 
  try:
  self._conn = mssql.connect(**self._args)
  except mssql.Error as e:
  raise ConnectError(*e.args) from e
 
  def quote(self, s: str):
- return f'[{s}]'
+ return f"[{s}]"
 
  def md5_to_int(self, s: str) -> str:
  return f"CONVERT(decimal(38,0), CONVERT(bigint, HashBytes('MD5', {s}), 2))"
@@ -188,13 +202,15 @@ def md5_to_int(self, s: str) -> str:
  def to_string(self, s: str):
  return f"CONVERT(varchar, {s})"
 
+
 class BigQuery(Database):
  def __init__(self, project, dataset):
  from google.cloud import bigquery
+
  self._client = bigquery.Client(project)
 
  def quote(self, s: str):
- return f'`{s}`'
+ return f"`{s}`"
 
  def md5_to_int(self, s: str) -> str:
  return f"cast(cast( ('0x' || substr(TO_HEX(md5({s})), 18)) as int64) as numeric)"
@@ -206,29 +222,27 @@ def _canonize_value(self, value):
 
  def _query(self, sql_code: str):
  from google.cloud import bigquery
+
  try:
  res = list(self._client.query(sql_code))
  except Exception as e:
  msg = "Exception when trying to execute SQL code:\n %s\n\nGot error: %s"
- raise ConnectError(msg%(sql_code, e))
+ raise ConnectError(msg % (sql_code, e))
 
  if res and isinstance(res[0], bigquery.table.Row):
  res = [tuple(self._canonize_value(v) for v in row.values()) for row in res]
  return res
 
  def to_string(self, s: str):
- return f'cast({s} as string)' 
+ return f"cast({s} as string)"
+
 
 class Snowflake(Database):
  def __init__(self, account, user, password, path, schema, database, print_sql=False):
  snowflake = import_snowflake()
- logging.getLogger('snowflake.connector').setLevel(logging.WARNING)
+ logging.getLogger("snowflake.connector").setLevel(logging.WARNING)
 
- self._conn = snowflake.connector.connect(
- user=user,
- password=password,
- account=account
- )
+ self._conn = snowflake.connector.connect(user=user, password=password, account=account)
  self._conn.cursor().execute(f"USE WAREHOUSE {path.lstrip('/')}")
  self._conn.cursor().execute(f"USE DATABASE {database}")
  self._conn.cursor().execute(f"USE SCHEMA {schema}")
@@ -240,7 +254,7 @@ def md5_to_int(self, s: str) -> str:
  return f"BITAND(md5_number_lower64({s}), {CHECKSUM_MASK})"
 
  def to_string(self, s: str):
- return f'cast({s} as string)' 
+ return f"cast({s} as string)"
 
 
 def connect_to_uri(db_uri: str) -> Database:
@@ -259,28 +273,28 @@ def connect_to_uri(db_uri: str) -> Database:
  dsn = dsnparse.parse(db_uri)
  if len(dsn.schemes) > 1:
  raise NotImplementedError("No support for multiple schemes")
- scheme ,= dsn.schemes
+ (scheme,) = dsn.schemes
 
  if len(dsn.paths) == 0:
- path = ''
+ path = ""
  elif len(dsn.paths) == 1:
- path ,= dsn.paths
+ (path,) = dsn.paths
  else:
  raise ValueError("Bad value for uri, too many paths: %s" % db_uri)
 
- if scheme == 'postgres':
+ if scheme == "postgres":
  return Postgres(dsn.host, dsn.port, path, dsn.user, dsn.password)
- elif scheme == 'mysql':
+ elif scheme == "mysql":
  return MySQL(dsn.host, dsn.port, path, dsn.user, dsn.password)
- elif scheme == 'snowflake':
+ elif scheme == "snowflake":
  return Snowflake(dsn.host, dsn.user, dsn.password, path, **dsn.query)
- elif scheme == 'mssql':
+ elif scheme == "mssql":
  return MsSQL(dsn.host, dsn.port, path, dsn.user, dsn.password)
- elif scheme == 'bigquery':
+ elif scheme == "bigquery":
  return BigQuery(dsn.host, path)
- elif scheme == 'redshift':
+ elif scheme == "redshift":
  return Redshift(dsn.host, dsn.port, path, dsn.user, dsn.password)
- elif scheme == 'oracle':
+ elif scheme == "oracle":
  return Oracle(dsn.host, dsn.port, path, dsn.user, dsn.password)
 
  raise NotImplementedError(f"Scheme {dsn.scheme} currently not supported")