Now accepting multiple -k arguments

Author: erezsh

data_diff/__main__.py

@@ -40,9 +40,9 @@ def _remove_passwords_in_dict(d: dict):
 @click.argument("table1", required=False)
 @click.argument("database2", required=False)
 @click.argument("table2", required=False)
-@click.option("-k", "--key-column", default=None, help="Name of primary key column. Default='id'.")
+@click.option("-k", "--key-column", multiple=True, help="Names of primary key columns. Default='id'.")
 @click.option("-t", "--update-column", default=None, help="Name of updated_at/last_updated column")
-@click.option("-c", "--columns", default=[], multiple=True, help="Names of extra columns to compare")
+@click.option("-c", "--columns", 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=None, help=f"Segments per iteration. Default={DEFAULT_BISECTION_FACTOR}.")
 @click.option(
@@ -130,7 +130,7 @@ def _main(
  logging.error("Cannot specify a limit when using the -s/--stats switch")
  return
 
- key_column = key_column or ("id",)
+ key_columns = tuple(key_column) or ("id",)
  if bisection_factor is None:
  bisection_factor = DEFAULT_BISECTION_FACTOR
  if bisection_threshold is None:
@@ -172,8 +172,8 @@ def _main(
  logging.error("Error while parsing age expression: %s" % e)
  return
 
- table1_seg = TableSegment(db1, db1.parse_table_name(table1), (key_column,), update_column, columns, **options)
- table2_seg = TableSegment(db2, db2.parse_table_name(table2), (key_column,), update_column, columns, **options)
+ table1_seg = TableSegment(db1, db1.parse_table_name(table1), key_columns, update_column, columns, **options)
+ table2_seg = TableSegment(db2, db2.parse_table_name(table2), key_columns, update_column, columns, **options)
 
  differ = TableDiffer(
  bisection_factor=bisection_factor,