Created
October 14, 2022 06:36
-
-
Save rik/c15479068812eef069951270e71e3e91 to your computer and use it in GitHub Desktop.
Slightly modified decorator from https://adamj.eu/tech/2022/10/13/dry-run-mode-for-data-imports-in-django/
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 argparse | |
import csv | |
from collections.abc import Generator | |
from contextlib import contextmanager, closing | |
from io import TextIOWrapper | |
from typing import Any | |
from django.core.management.base import BaseCommand | |
from django.db import transaction | |
from example.core.models import Author, Book | |
class Command(BaseCommand): | |
help = "Import books from a CSV" | |
def add_arguments(self, parser): | |
parser.add_argument( | |
"--write", | |
action="store_true", | |
default=False, | |
help="Actually edit the database", | |
) | |
parser.add_argument("file", type=argparse.FileType()) | |
def handle( | |
self, *args: Any, file: TextIOWrapper, write: bool, **kwargs: Any | |
) -> None: | |
if not write: | |
self.stdout.write("In dry run mode (--write not passed)") | |
books_created = 0 | |
authors_created = 0 | |
with closing(file), atomic(commit=write): | |
csvfile = csv.reader(file) | |
header = next(csvfile) | |
if header != ["Title", "Author"]: | |
self.stdout.write("Unexpected header row, should be “Title,Author”") | |
raise SystemExit(1) | |
for title, author_name in csvfile: | |
author, created = Author.objects.get_or_create(name=author_name) | |
if created: | |
authors_created += 1 | |
book, created = Book.objects.get_or_create( | |
title=title, | |
author=author, | |
) | |
if created: | |
books_created += 1 | |
if write: | |
prefix = "Created" | |
else: | |
prefix = "Would create" | |
self.stdout.write( | |
f"{prefix} {books_created} books and {authors_created} authors." | |
) | |
class DoRollback(Exception): | |
pass | |
@contextmanager | |
def atomic(*, commit: bool) -> Generator[None, None, None]: | |
try: | |
with transaction.atomic(): | |
yield | |
if not commit: | |
raise DoRollback() | |
except DoRollback: | |
pass |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment