Created
August 15, 2019 12:42
-
-
Save Sinkler/0870dd466aa46498c843c6071f6c3a68 to your computer and use it in GitHub Desktop.
Django Postgres Copy From
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import csv | |
from io import StringIO | |
from django.db import connection | |
from psycopg2 import sql | |
def copy_from(model, columns, records): | |
""" | |
copy_from( | |
UserFeed, | |
['user_id', 'feed_id', 'sources', 'new', 'created', 'hidden'], | |
records | |
) | |
""" | |
stream = StringIO() | |
writer = csv.writer(stream) | |
for record in records: | |
writer.writerow(record) | |
stream.seek(0) | |
with connection.cursor() as cursor: | |
cursor.copy_from( | |
file=stream, | |
table=model._meta.db_table, | |
sep=',', | |
columns=columns, | |
) | |
def copy_from_conflict(model, columns, records, on_conflict, tmp_name): | |
""" | |
copy_from_conflict( | |
UserFeed, | |
['user_id', 'feed_id', 'sources', 'new', 'created', 'hidden'], | |
records, | |
"(user_id, feed_id) DO UPDATE SET sources = ARRAY(SELECT DISTINCT UNNEST(a.sources || EXCLUDED.sources))", | |
'temp_user_feed_game_{}'.format(feed.id) | |
) | |
""" | |
stream = StringIO() | |
writer = csv.writer(stream) | |
for record in records: | |
writer.writerow(record) | |
stream.seek(0) | |
columns = sql.SQL(', ').join([sql.Identifier(c) for c in columns]) | |
tmp_name = sql.Identifier(tmp_name) | |
name = sql.Identifier(model._meta.db_table) | |
tmp = sql.SQL('CREATE TEMP TABLE {} ON COMMIT DROP AS SELECT * FROM {} WITH NO DATA').format(tmp_name, name) | |
cp = sql.SQL('COPY {} ({}) FROM STDIN WITH CSV').format(tmp_name, columns) | |
ins = sql.SQL('INSERT INTO {0} AS a ({1}) SELECT {1} FROM {2} b ON CONFLICT {3}').format( | |
name, | |
columns, | |
tmp_name, | |
sql.SQL(on_conflict) | |
) | |
with connection.cursor() as cursor: | |
cursor.execute(tmp) | |
cursor.copy_expert(cp, stream) | |
cursor.execute(ins) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment