bgneal@292: """
bgneal@292: import_old_news_comments.py - For importing comments on news stories from SG101 1.0 as csv files.
bgneal@292: """
bgneal@292: from __future__ import with_statement
bgneal@292: import csv
bgneal@292: import optparse
bgneal@292: import sys
bgneal@292: from datetime import datetime
bgneal@292: 
bgneal@292: from django.core.management.base import LabelCommand, CommandError
bgneal@292: from django.contrib.auth.models import User
bgneal@292: from django.contrib.contenttypes.models import ContentType
bgneal@292: 
bgneal@292: from comments.models import Comment
bgneal@292: from news.models import Story
bgneal@292: import legacy.data
bgneal@292: from legacy.html2md import MarkdownWriter
bgneal@292: 
bgneal@292: 
bgneal@292: class Command(LabelCommand):
bgneal@292:     args = '<filename filename ...>'
bgneal@292:     help = 'Imports news story comments from the old database in CSV format'
bgneal@292:     option_list = LabelCommand.option_list + (
bgneal@292:         optparse.make_option("-p", "--progress", action="store_true",
bgneal@292:             help="Output a . after every 20 comments to show progress"),
bgneal@292:     )
bgneal@292:     md_writer = MarkdownWriter()
bgneal@292: 
bgneal@292:     def handle_label(self, filename, **options):
bgneal@292:         """
bgneal@292:         Process each line in the CSV file given by filename by
bgneal@292:         creating a new story comment.
bgneal@292: 
bgneal@292:         """
bgneal@292:         self.show_progress = options.get('progress')
bgneal@292:         self.users = {}
bgneal@292: 
bgneal@292:         try:
bgneal@292:             with open(filename, "rb") as f:
bgneal@292:                 self.reader = csv.DictReader(f)
bgneal@292:                 num_rows = 0
bgneal@292:                 try:
bgneal@292:                     for row in self.reader:
bgneal@292:                         self.process_row(row)
bgneal@292:                         num_rows += 1
bgneal@292:                         if self.show_progress and num_rows % 20 == 0:
bgneal@292:                             sys.stdout.write('.')
bgneal@292:                             sys.stdout.flush()
bgneal@292:                 except csv.Error, e:
bgneal@292:                     raise CommandError("CSV error: %s %s %s" % (
bgneal@292:                         filename, self.reader.line_num, e))
bgneal@292: 
bgneal@292:                 print
bgneal@292: 
bgneal@292:         except IOError:
bgneal@292:             raise CommandError("Could not open file: %s" % filename)
bgneal@292: 
bgneal@292:     def process_row(self, row):
bgneal@292:         """
bgneal@292:         Process one row from the CSV file: create a Comment object for
bgneal@292:         the row and save it in the database.
bgneal@292: 
bgneal@292:         """
bgneal@292:         row = dict((k, v if v != 'NULL' else '') for k, v in row.iteritems())
bgneal@292: 
bgneal@292:         try:
bgneal@292:             user = self._get_user(row['name'])
bgneal@292:         except User.DoesNotExist:
bgneal@292:             print "Could not find user %s for comment %s; skipping." % (
bgneal@292:                     row['name'], row['tid'])
bgneal@292:             return
bgneal@292: 
bgneal@292:         try:
bgneal@292:             story = Story.objects.get(id=int(row['sid']))
bgneal@292:         except Story.DoesNotExist:
bgneal@292:             print "Could not find story %s for comment %s; skipping." % (
bgneal@292:                     row['sid'], row['tid'])
bgneal@292:             return
bgneal@292: 
bgneal@292:         comment = Comment(
bgneal@292:             id=int(row['tid']),
bgneal@292:             content_type = ContentType.objects.get_for_model(story),
bgneal@292:             object_id = story.id,
bgneal@292:             user = user,
bgneal@292:             comment = self.to_markdown(row['comment']),
bgneal@292:             creation_date = datetime.strptime(row['date'], "%Y-%m-%d %H:%M:%S"),
bgneal@292:             ip_address = row['host_name'],
bgneal@292:             is_public = True,
bgneal@292:             is_removed = False,
bgneal@292:         )
bgneal@292: 
bgneal@292:         comment.save()
bgneal@292: 
bgneal@292:     def _get_user(self, username):
bgneal@292:         """
bgneal@292:         Returns the user object with the given username.
bgneal@292:         Throws User.DoesNotExist if not found.
bgneal@292: 
bgneal@292:         """
bgneal@292:         try:
bgneal@292:             return self.users[username]
bgneal@292:         except KeyError:
bgneal@292:             pass
bgneal@292: 
bgneal@292:         try:
bgneal@292:             user = User.objects.get(username=username)
bgneal@292:         except User.DoesNotExist:
bgneal@292:             try:
bgneal@292:                 user = User.objects.get(
bgneal@292:                         username=legacy.data.KNOWN_USERNAME_CHANGES[username])
bgneal@292:             except KeyError:
bgneal@292:                 raise User.DoesNotExist
bgneal@292: 
bgneal@292:         self.users[username] = user
bgneal@292:         return user
bgneal@292: 
bgneal@292:     def to_markdown(self, s):
bgneal@292:         self.md_writer.reset()
bgneal@292: 
bgneal@292:         if not isinstance(s, unicode):
bgneal@294:             s = s.decode('latin-1', 'replace')
bgneal@292: 
bgneal@292:         self.md_writer.feed(s)
bgneal@292:         return self.md_writer.markdown()