import_message_data: Use psycopg2.extras.execute_values.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
This commit is contained in:
Anders Kaseorg 2020-06-09 02:01:54 -07:00 committed by Tim Abbott
parent 6649c8e35f
commit 2604ebba38
2 changed files with 11 additions and 6 deletions

View File

@ -34,7 +34,7 @@ class TimeTrackingCursor(cursor):
return wrapper_execute(self, super().execute, query, vars) return wrapper_execute(self, super().execute, query, vars)
def executemany(self, query: Query, def executemany(self, query: Query,
vars: Iterable[Params]) -> 'TimeTrackingCursor': vars: Iterable[Params]) -> 'TimeTrackingCursor': # nocoverage
return wrapper_execute(self, super().executemany, query, vars) return wrapper_execute(self, super().executemany, query, vars)
class TimeTrackingConnection(connection): class TimeTrackingConnection(connection):

View File

@ -12,6 +12,8 @@ from django.db.models import Max
from django.utils.timezone import now as timezone_now from django.utils.timezone import now as timezone_now
from typing import Any, Dict, List, Optional, Set, Tuple, \ from typing import Any, Dict, List, Optional, Set, Tuple, \
Iterable, cast Iterable, cast
from psycopg2.extras import execute_values
from psycopg2.sql import Identifier, SQL
from analytics.models import RealmCount, StreamCount, UserCount from analytics.models import RealmCount, StreamCount, UserCount
from zerver.lib.actions import UserMessageLite, bulk_insert_ums, \ from zerver.lib.actions import UserMessageLite, bulk_insert_ums, \
@ -1289,12 +1291,15 @@ def import_attachments(data: TableData) -> None:
# TODO: Do this the kosher Django way. We may find a # TODO: Do this the kosher Django way. We may find a
# better way to do this in Django 1.9 particularly. # better way to do this in Django 1.9 particularly.
with connection.cursor() as cursor: with connection.cursor() as cursor:
sql_template = ''' sql_template = SQL('''
insert into %s (%s, %s) values(%%s, %%s);''' % (m2m_table_name, insert into {} ({}, {}) values %s
parent_id, ''').format(
child_id) Identifier(m2m_table_name),
Identifier(parent_id),
Identifier(child_id),
)
tups = [(row[parent_id], row[child_id]) for row in m2m_rows] tups = [(row[parent_id], row[child_id]) for row in m2m_rows]
cursor.executemany(sql_template, tups) execute_values(cursor.cursor, sql_template, tups)
logging.info('Successfully imported M2M table %s', m2m_table_name) logging.info('Successfully imported M2M table %s', m2m_table_name)