/mysql2pgsql/lib/postgres_file_writer.py

https://github.com/foxxtrot/py-mysql2pgsql · Python · 157 lines · 93 code · 16 blank · 48 comment · 10 complexity · ece981d5710fee64854917db5a90a929 MD5 · raw file

  1. from __future__ import absolute_import
  2. import time
  3. import sys
  4. from cStringIO import StringIO
  5. from psycopg2.extensions import QuotedString
  6. from .postgres_writer import PostgresWriter
  7. from . import print_row_progress, status_logger
  8. class PostgresFileWriter(PostgresWriter):
  9. """Class used to ouput the PostgreSQL
  10. compatable DDL and/or data to the specified
  11. output :py:obj:`file` from a MySQL server.
  12. :Parameters:
  13. - `output_file`: the output :py:obj:`file` to send the DDL and/or data
  14. - `verbose`: whether or not to log progress to :py:obj:`stdout`
  15. """
  16. verbose = None
  17. def __init__(self, output_file, verbose=False):
  18. self.verbose = verbose
  19. self.f = output_file
  20. self.f.write("""
  21. -- MySQL 2 PostgreSQL dump\n
  22. SET client_encoding = 'UTF8';
  23. SET standard_conforming_strings = off;
  24. SET check_function_bodies = false;
  25. SET client_min_messages = warning;
  26. """)
  27. @status_logger
  28. def truncate(self, table):
  29. """Write DDL to truncate the specified `table`
  30. :Parameters:
  31. - `table`: an instance of a :py:class:`mysql2pgsql.lib.mysql_reader.MysqlReader.Table` object that represents the table to read/write.
  32. Returns None
  33. """
  34. truncate_sql, serial_key_sql = super(self.__class__, self).truncate(table)
  35. self.f.write("""
  36. -- TRUNCATE %(table_name)s;
  37. %(truncate_sql)s
  38. """ % {'table_name': table.name, 'truncate_sql': truncate_sql})
  39. if serial_key_sql:
  40. self.f.write("""
  41. %(serial_key_sql)s
  42. """ % {
  43. 'serial_key_sql': serial_key_sql})
  44. @status_logger
  45. def write_table(self, table):
  46. """Write DDL to create the specified `table`.
  47. :Parameters:
  48. - `table`: an instance of a :py:class:`mysql2pgsql.lib.mysql_reader.MysqlReader.Table` object that represents the table to read/write.
  49. Returns None
  50. """
  51. table_sql, serial_key_sql = super(self.__class__, self).write_table(table)
  52. if serial_key_sql:
  53. self.f.write("""
  54. %(serial_key_sql)s
  55. """ % {
  56. 'serial_key_sql': '\n'.join(serial_key_sql)
  57. })
  58. self.f.write("""
  59. -- Table: %(table_name)s
  60. %(table_sql)s
  61. """ % {
  62. 'table_name': table.name,
  63. 'table_sql': '\n'.join(table_sql),
  64. })
  65. @status_logger
  66. def write_indexes(self, table):
  67. """Write DDL of `table` indexes to the output file
  68. :Parameters:
  69. - `table`: an instance of a :py:class:`mysql2pgsql.lib.mysql_reader.MysqlReader.Table` object that represents the table to read/write.
  70. Returns None
  71. """
  72. self.f.write('\n'.join(super(self.__class__, self).write_indexes(table)))
  73. @status_logger
  74. def write_constraints(self, table):
  75. """Write DDL of `table` constraints to the output file
  76. :Parameters:
  77. - `table`: an instance of a :py:class:`mysql2pgsql.lib.mysql_reader.MysqlReader.Table` object that represents the table to read/write.
  78. Returns None
  79. """
  80. self.f.write('\n'.join(super(self.__class__, self).write_constraints(table)))
  81. @status_logger
  82. def write_contents(self, table, reader):
  83. """Write the data contents of `table` to the output file.
  84. :Parameters:
  85. - `table`: an instance of a :py:class:`mysql2pgsql.lib.mysql_reader.MysqlReader.Table` object that represents the table to read/write.
  86. - `reader`: an instance of a :py:class:`mysql2pgsql.lib.mysql_reader.MysqlReader` object that allows reading from the data source.
  87. Returns None
  88. """
  89. # start variable optimiztions
  90. pr = self.process_row
  91. f_write = self.f.write
  92. verbose = self.verbose
  93. # end variable optimiztions
  94. f_write("""
  95. --
  96. -- Data for Name: %(table_name)s; Type: TABLE DATA;
  97. --
  98. COPY "%(table_name)s" (%(column_names)s) FROM stdin;
  99. """ % {
  100. 'table_name': table.name,
  101. 'column_names': ', '.join(('"%s"' % col['name']) for col in table.columns)})
  102. if verbose:
  103. tt = time.time
  104. start_time = tt()
  105. prev_val_len = 0
  106. prev_row_count = 0
  107. for i, row in enumerate(reader.read(table)):
  108. row = list(row)
  109. pr(table, row)
  110. try:
  111. f_write('%s\n' % ('\t'.join(row)))
  112. except UnicodeDecodeError:
  113. f_write('%s\n' % ('\t'.join(row)).decode('utf-8'))
  114. if verbose:
  115. if ((i + 1) % 20000) == 0:
  116. now = tt()
  117. elapsed = now - start_time
  118. val = '%.2f rows/sec [%s] ' % (((i + 1) - prev_row_count) / elapsed, (i + 1))
  119. print_row_progress('%s%s' % (("\b" * prev_val_len), val))
  120. prev_val_len = len(val) + 3
  121. start_time = now
  122. prev_row_count = i + 1
  123. f_write("\\.\n\n")
  124. if verbose:
  125. print('')
  126. def close(self):
  127. """Closes the output :py:obj:`file`"""
  128. self.f.close()