Merge pull request #198 from coursera/cliu_dont_pass_in_command_twice

dont pass in command twice
coursera · Dec 3, 2015 · 011d02e · 011d02e
2 parents 749d0ac + b4b3fc3
commit 011d02e
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 6 deletions.
diff --git a/dataduct/steps/column_check.py b/dataduct/steps/column_check.py
@@ -23,7 +23,7 @@ def __init__(self, id, source_sql, source_host,
                  destination_table_definition=None, script=None,
                  destination_sql=None, sql_tail_for_source=None,
                  sample_size=100, tolerance=1.0, script_arguments=None,
-                 log_to_s3=False, **kwargs):
+                 log_to_s3=False, command=None, **kwargs):
         """Constructor for the ColumnCheckStep class
 
         Args:
@@ -60,7 +60,8 @@ def __init__(self, id, source_sql, source_host,
         if log_to_s3:
             script_arguments.append('--log_to_s3')
 
-        command = None if script else const.COLUMN_CHECK_COMMAND
+        if script is None and command is None:
+            command = const.COLUMN_CHECK_COMMAND
 
         super(ColumnCheckStep, self).__init__(
             id=id, script=script, command=command,

diff --git a/dataduct/steps/count_check.py b/dataduct/steps/count_check.py
@@ -22,7 +22,7 @@ def __init__(self, id, source_host, source_sql=None,
                  source_table_name=None, destination_table_name=None,
                  destination_table_definition=None, destination_sql=None,
                  tolerance=1.0, script_arguments=None, log_to_s3=False,
-                 script=None, source_count_sql=None, **kwargs):
+                 script=None, source_count_sql=None, command=None, **kwargs):
         """Constructor for the CountCheckStep class
 
         Args:
@@ -66,7 +66,8 @@ def __init__(self, id, source_host, source_sql=None,
         if log_to_s3:
             script_arguments.append('--log_to_s3')
 
-        command = None if script else const.COUNT_CHECK_COMMAND
+        if script is None and command is None:
+            command = const.COUNT_CHECK_COMMAND
 
         super(CountCheckStep, self).__init__(
             id=id, command=command, script=script,

diff --git a/dataduct/steps/primary_key_check.py b/dataduct/steps/primary_key_check.py
@@ -16,7 +16,7 @@ class PrimaryKeyCheckStep(QATransformStep):
     """
 
     def __init__(self, id, table_definition, script_arguments=None,
-                 log_to_s3=False, **kwargs):
+                 log_to_s3=False, command=None, script=None, **kwargs):
         """Constructor for the PrimaryKeyCheckStep class
 
         Args:
@@ -36,6 +36,9 @@ def __init__(self, id, table_definition, script_arguments=None,
         if log_to_s3:
             script_arguments.append('--log_to_s3')
 
+        if script is None and command is None:
+            command = const.PK_CHECK_COMMAND
+
         super(PrimaryKeyCheckStep, self).__init__(
-            id=id, command=const.PK_CHECK_COMMAND,
+            id=id, command=command, script=script,
             script_arguments=script_arguments, **kwargs)