最有效的psycopg2 Python转义函数

Question

我现在正在尝试使用psycopg2库的(copy_from)功能进行批量插入。下面是我的代码。我希望能得到一些关于最有效的转义函数和处理包含限制字符的值的最佳实践的建议。我使用的是python3，字符串来自mongodb。

这些限制字符可以在这里找到： http://www.postgresql.org/docs/9.2/static/sql-copy.html 文档中提到： “反斜杠字符（\）可以在COPY数据中用来引用那些可能被误认为是行或列分隔符的数据字符。特别是，如果以下字符作为列值的一部分出现，它们必须前面加上反斜杠：反斜杠本身、换行符、回车符，以及当前的分隔符字符。”

def bulk_write(self, table, data, columns = None):
    with psycopg2.connect(database = self.database,
                          user = self.user,
                          host = self.host,
                          password = self.password) as conn:
      with conn.cursor() as cur:
        cur.execute("SET TIME ZONE 'PDT8PST';")
        cols_import = tuple(columns) if columns else None
        data_tsv = '\n'.join(['\t'.join(self.escape_copy_string(str(value)) for value in row) for row in data])
        with open("test","w") as f:
          f.write(data_tsv)
        cur.copy_from(io.StringIO(data_tsv), table, columns=cols_import, null="None")

   def escape_copy_string(self,s):
      s = s.replace("\\","\\\\").replace("\n","\\n").replace("\r","\\r").replace("\t","\\t")
      return s

Answer 1

经过一些测试，我觉得使用 replace 方法是最有效的，至少和基于正则表达式的方法相比是这样：

import re
import timeit

string_rep = {
    '\\': '\\\\',
    '\r': '\\r',
    '\n': '\\n',
    '\t': '\\t',
}
string_pattern = re.compile('|'.join(re.escape(key) for key in string_rep))


def escape_re_sub(text):
    return string_pattern.sub(lambda m: string_rep[m.group(0)], text)


def escape_str_replace(text):
    return (
        text.replace('\\', '\\\\')
        .replace('\n', '\\n')
        .replace('\r', '\\r')
        .replace('\t', '\\t')
    )


string_no_escape = 'This is some bit of text that has no strings to replace'

time_re_sub = timeit.Timer(lambda: escape_re_sub(string_no_escape)).autorange()
print('no escape sub    ', time_re_sub[0] / time_re_sub[1], 'iterations per second')

time_str_replace = timeit.Timer(lambda: escape_str_replace(string_no_escape)).autorange()
print('no escape replace', time_str_replace[0] / time_str_replace[1], 'iterations per second')

string_escape = 'This is some\r \\ bit of text \nthat has \t\t some things to \\t replace'

time_re_sub = timeit.Timer(lambda: escape_re_sub(string_escape)).autorange()
print('with escape sub    ', time_re_sub[0] / time_re_sub[1], 'iterations per second')

time_str_replace = timeit.Timer(lambda: escape_str_replace(string_escape)).autorange()
print('with escape replace', time_str_replace[0] / time_str_replace[1], 'iterations per second')

对我来说，输出结果是：

no escape sub     1088292.3082792824 iterations per second
no escape replace 1310796.652683603 iterations per second
with escape sub     251530.53121228397 iterations per second
with escape replace 913308.513839589 iterations per second

当需要转义的字符存在时，这种差别尤其明显，但即使在转义没有改变字符串的情况下，这种差别也是存在的。

Answer 2

与其自己手动去写，不如使用 csv 模块，并在 csv 模式下使用 copy_from。

如果你用列表推导式来处理数据，那数据应该不会太大。因为如果数据量大，你很快就会耗尽内存。可以考虑用一个循环，逐行写入数据，这样会更好。