Merge branch 'master' of github.com:Oefenweb/python-untraceables

Oefenweb · Apr 7, 2017 · c9fc530 · c9fc530
2 parents 692ddaf + 7be058f
commit c9fc530
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 1 deletion.
diff --git a/bin/randomize-ids b/bin/randomize-ids
@@ -155,6 +155,7 @@ def run_sql(args):
 
   database = args.database
   foreign_key_checks = args.foreign_key_checks
+  unique_checks = args.unique_checks
 
   if sys.stdin.isatty():
     cli_utility.print_e('Could not read any data from stdin')
@@ -165,6 +166,12 @@ def run_sql(args):
     foreign_key_checks_off = [query_utility.get_foreign_key_checks(foreign_key_checks)]
     foreign_key_checks_on = [query_utility.get_foreign_key_checks(not foreign_key_checks)]
 
+  unique_checks_off = []
+  unique_checks_on = []
+  if not foreign_key_checks:
+    unique_checks_off = [query_utility.get_unique_checks(unique_checks)]
+    unique_checks_on = [query_utility.get_unique_checks(not unique_checks)]
+
   statements_from_stdin = mysql_utility.split_file(sys.stdin)
 
   connection = cursor = None
@@ -173,7 +180,9 @@ def run_sql(args):
     connection.autocommit(True)
     cursor = mysql_utility.get_cursor(connection)
 
-    statements = chain(iter(foreign_key_checks_off), statements_from_stdin, iter(foreign_key_checks_on))
+    statements = chain(iter(foreign_key_checks_off), iter(unique_checks_off),
+                       statements_from_stdin,
+                       iter(unique_checks_on), iter(foreign_key_checks_on))
     for statement in statements:
       stripped_statement = statement.strip()
       if stripped_statement != '':
@@ -272,6 +281,15 @@ def main():
                               dest='foreign_key_checks',
                               help='Whether or not to enable FOREIGN_KEY_CHECKS')
   run_sql_parser.set_defaults(foreign_key_checks=True)
+  run_sql_parser.add_argument('--unique-checks',
+                              action='store_true',
+                              dest='unique_checks',
+                              help='Whether or not to enable UNIQUE_CHECKS')
+  run_sql_parser.add_argument('--no-unique-checks',
+                              action='store_false',
+                              dest='unique_checks',
+                              help='Whether or not to enable UNIQUE_CHECKS')
+  run_sql_parser.set_defaults(foreign_key_checks=True)
 
   parser.add_argument('-v', '--verbose', action='store_true', help='Be more verbose')
 

diff --git a/untraceables/test/test_query.py b/untraceables/test/test_query.py
@@ -63,6 +63,23 @@ def test_get_foreign_key_checks(self):
     actual = query_utility.get_foreign_key_checks(0)
     self.assertEqual(expected, actual)
 
+  def test_get_unique_checks(self):
+    """
+    Tests `get_unique_checks`.
+    """
+
+    actual = query_utility.get_unique_checks(True)
+    expected = 'SET UNIQUE_CHECKS=1'
+    self.assertEqual(expected, actual)
+    actual = query_utility.get_unique_checks(1)
+    self.assertEqual(expected, actual)
+
+    actual = query_utility.get_unique_checks(False)
+    expected = 'SET UNIQUE_CHECKS=0'
+    self.assertEqual(expected, actual)
+    actual = query_utility.get_unique_checks(0)
+    self.assertEqual(expected, actual)
+
   def test_get_randomize(self):
     """
     Tests `get_randomize`.

diff --git a/untraceables/utilities/query.py b/untraceables/utilities/query.py
@@ -72,6 +72,19 @@ def get_foreign_key_checks(enabled):
   return 'SET FOREIGN_KEY_CHECKS={0:d}'.format(enabled)
 
 
+def get_unique_checks(enabled):
+  """
+  Gets the query the enable / disable UNIQUE_CHECKS.
+
+  :type bool
+  :param enabled: Whether or not to enable
+  :rtype str
+  :return A query
+  """
+
+  return 'SET UNIQUE_CHECKS={0:d}'.format(enabled)
+
+
 def get_randomize(database, table, columns, column, mapping_database, mapping_table):
   """
   Gets the queries to randomize a table / column in a given database.