projects.mako.cc - protection-tools/blob - 02-mysqldump_to_csv.py

   1 #!/usr/bin/env python
   2 import fileinput
   3 import csv
   4 import sys
   5 import re
   6
   7 # This prevents prematurely closed pipes from raising
   8 # an exception in Python
   9 from signal import signal, SIGPIPE, SIG_DFL
  10 signal(SIGPIPE, SIG_DFL)
  11
  12 def is_insert(line):
  13     """
  14     Returns true if the line begins a SQL insert statement.
  15     """
  16     return line.startswith('INSERT INTO') or False
  17
  18
  19 def get_values(line):
  20     """
  21     Returns the portion of an INSERT statement containing values
  22     """
  23     return line.partition('` VALUES ')[2].strip()
  24
  25
  26 def values_sanity_check(values):
  27     """
  28     Ensures that values from the INSERT statement meet basic checks.
  29     """
  30     assert values
  31     assert values[0] == '('
  32     assert values[-2:] == ');'
  33     # Assertions have not been raised
  34     return True
  35
  36 def clean_quotes_for_fread(col):
  37     col = re.sub('\t', ' ', col)
  38     if col.startswith('"'):
  39         return('"' + col + '"')
  40     else:
  41         return(col)
  42
  43 def parse_values(values, outfile):
  44     """
  45     Given a file handle and the raw values from a MySQL INSERT
  46     statement, write the equivalent CSV to the file
  47     """
  48     values = values.rstrip(");")
  49     values = values.lstrip("(")
  50
  51     reader = csv.reader(values.split("),("), delimiter=',',
  52                         doublequote=False,
  53                         escapechar='\\',
  54                         quotechar="'",
  55                         strict=True)
  56
  57     for reader_row in reader:
  58         print("\t".join([clean_quotes_for_fread(col) for col in reader_row]))
  59
  60 def main():
  61     """
  62     Parse arguments and start the program
  63     """
  64     # Iterate over all lines in all files
  65     # listed in sys.argv[1:]
  66     # or stdin if no args given.
  67     try:
  68         for line in fileinput.input():
  69             # Look for an INSERT statement and parse it.
  70             if is_insert(line):
  71                 values = get_values(line.strip().rstrip())
  72                 if values_sanity_check(values):
  73                     parse_values(values, sys.stdout)
  74     except KeyboardInterrupt:
  75         sys.exit(0)
  76
  77 if __name__ == "__main__":
  78     main()