bgneal@294
|
1 """
|
bgneal@294
|
2 import_old_links.py - For importing links from SG101 1.0 as csv files.
|
bgneal@294
|
3 """
|
bgneal@333
|
4 from __future__ import with_statement
|
bgneal@294
|
5 import csv
|
bgneal@294
|
6 import datetime
|
bgneal@294
|
7
|
bgneal@294
|
8 from django.core.management.base import LabelCommand, CommandError
|
bgneal@294
|
9 from django.contrib.auth.models import User
|
bgneal@294
|
10
|
bgneal@294
|
11 from weblinks.models import Link, Category
|
bgneal@294
|
12
|
bgneal@294
|
13
|
bgneal@294
|
14 class Command(LabelCommand):
|
bgneal@294
|
15 args = '<filename filename ...>'
|
bgneal@294
|
16 help = 'Imports weblinks from the old database in CSV format'
|
bgneal@294
|
17
|
bgneal@294
|
18 def handle_label(self, filename, **options):
|
bgneal@294
|
19 """
|
bgneal@294
|
20 Process each line in the CSV file given by filename by
|
bgneal@294
|
21 creating a new weblink object and saving it to the database.
|
bgneal@294
|
22
|
bgneal@294
|
23 """
|
bgneal@294
|
24 self.cats = {}
|
bgneal@294
|
25 try:
|
bgneal@294
|
26 self.default_user = User.objects.get(pk=2)
|
bgneal@294
|
27 except User.DoesNotExist:
|
bgneal@294
|
28 raise CommandError("Need a default user with pk=1")
|
bgneal@294
|
29
|
bgneal@294
|
30 try:
|
bgneal@294
|
31 with open(filename, "rb") as f:
|
bgneal@294
|
32 self.reader = csv.DictReader(f)
|
bgneal@294
|
33 try:
|
bgneal@294
|
34 for row in self.reader:
|
bgneal@294
|
35 self.process_row(row)
|
bgneal@294
|
36 except csv.Error, e:
|
bgneal@294
|
37 raise CommandError("CSV error: %s %s %s" % (
|
bgneal@294
|
38 filename, self.reader.line_num, e))
|
bgneal@294
|
39
|
bgneal@294
|
40 except IOError:
|
bgneal@294
|
41 raise CommandError("Could not open file: %s" % filename)
|
bgneal@294
|
42
|
bgneal@294
|
43 def get_category(self, row):
|
bgneal@294
|
44 """
|
bgneal@294
|
45 Return the Category object for the row.
|
bgneal@294
|
46
|
bgneal@294
|
47 """
|
bgneal@294
|
48 cat_id = row['cid']
|
bgneal@294
|
49 if cat_id not in self.cats:
|
bgneal@294
|
50 try:
|
bgneal@294
|
51 cat = Category.objects.get(pk=cat_id)
|
bgneal@294
|
52 except Category.DoesNotExist:
|
bgneal@294
|
53 raise CommandError("Category does not exist: %s on line %s" % (
|
bgneal@294
|
54 cat_id, self.reader.line_num))
|
bgneal@294
|
55 else:
|
bgneal@294
|
56 self.cats[cat_id] = cat
|
bgneal@294
|
57 return self.cats[cat_id]
|
bgneal@294
|
58
|
bgneal@294
|
59 def get_user(self, username):
|
bgneal@294
|
60 """
|
bgneal@294
|
61 Return the user object for the given username.
|
bgneal@294
|
62 If the user cannot be found, self.default_user is returned.
|
bgneal@294
|
63
|
bgneal@294
|
64 """
|
bgneal@294
|
65 try:
|
bgneal@294
|
66 return User.objects.get(username=username)
|
bgneal@294
|
67 except User.DoesNotExist:
|
bgneal@294
|
68 return self.default_user
|
bgneal@294
|
69
|
bgneal@294
|
70 def process_row(self, row):
|
bgneal@294
|
71 """
|
bgneal@294
|
72 Process one row from the CSV file: create an object for the row
|
bgneal@294
|
73 and save it in the database.
|
bgneal@294
|
74
|
bgneal@294
|
75 """
|
bgneal@294
|
76 link = Link(category=self.get_category(row),
|
bgneal@294
|
77 title=row['title'].decode('latin-1'),
|
bgneal@294
|
78 url=row['url'].decode('latin-1'),
|
bgneal@294
|
79 description=row['description'].decode('latin-1'),
|
bgneal@294
|
80 user=self.get_user(row['submitter']),
|
bgneal@294
|
81 date_added=datetime.datetime.strptime(row['date'], "%Y-%m-%d %H:%M:%S"),
|
bgneal@294
|
82 hits=int(row['hits']),
|
bgneal@294
|
83 is_public=True)
|
bgneal@294
|
84 link.save()
|