forum_modules/exporter/importer.py

   1 import os, tarfile, datetime, ConfigParser, logging
   2
   3 from django.utils.translation import ugettext as _
   4 from django.core.cache import cache
   5
   6 from south.db import db
   7
   8 from xml.sax import make_parser
   9 from xml.sax.handler import ContentHandler, ErrorHandler
  10
  11 from forum.templatetags.extra_tags import diff_date
  12
  13 from exporter import TMP_FOLDER, DATETIME_FORMAT, DATE_FORMAT, META_INF_SECTION, CACHE_KEY
  14 from orm import orm
  15 import commands
  16
  17 NO_DEFAULT = object()
  18
  19 class ContentElement():
  20     def __init__(self, content):
  21         self._content = content
  22
  23     def content(self):
  24         return self._content.strip()
  25
  26     def as_bool(self):
  27         return self.content() == "true"
  28
  29     def as_date(self, default=NO_DEFAULT):
  30         try:
  31             return datetime.datetime.strptime(self.content(), DATE_FORMAT)
  32         except:
  33             if default == NO_DEFAULT:
  34                 return datetime.date.fromtimestamp(0)
  35             else:
  36                 return default
  37
  38
  39     def as_datetime(self, default=NO_DEFAULT):
  40         try:
  41             return datetime.datetime.strptime(self.content(), DATETIME_FORMAT)
  42         except:
  43             if default == NO_DEFAULT:
  44                 return datetime.datetime.fromtimestamp(0)
  45             else:
  46                 return default
  47
  48     def as_int(self, default=0):
  49         try:
  50             return int(self.content())
  51         except:
  52             return default
  53
  54     def __str__(self):
  55         return self.content()
  56
  57
  58 class RowElement(ContentElement):
  59     def __init__(self, name, attrs, parent=None):
  60         self.name = name.lower()
  61         self.parent = parent
  62         self.attrs = dict([(k.lower(), ContentElement(v)) for k, v in attrs.items()])
  63         self._content = u''
  64         self.sub_elements = {}
  65
  66         if parent:
  67             parent.add(self)
  68
  69     def add_to_content(self, ch):
  70         self._content += unicode(ch)
  71
  72     def add(self, sub):
  73         curr = self.sub_elements.get(sub.name, None)
  74
  75         if not curr:
  76             curr = []
  77             self.sub_elements[sub.name] = curr
  78
  79         curr.append(sub)
  80
  81     def get(self, name, default=None):
  82         return self.sub_elements.get(name.lower(), [default])[-1]
  83
  84     def get_list(self, name):
  85         return self.sub_elements.get(name.lower(), [])
  86
  87     def get_listc(self, name):
  88         return [r.content() for r in self.get_list(name)]
  89
  90     def getc(self, name, default=""):
  91         el = self.get(name, None)
  92
  93         if el:
  94             return el.content()
  95         else:
  96             return default
  97
  98     def get_attr(self, name, default=""):
  99         return self.attrs.get(name.lower(), default)
 100
 101     def as_pickled(self, default=None):
 102         value_el = self.get('value')
 103
 104         if value_el:
 105             return value_el._as_pickled(default)
 106         else:
 107             return default
 108
 109     TYPES_MAP = dict([(c.__name__, c) for c in (int, long, str, unicode, float)])
 110
 111     def _as_pickled(self, default=None):
 112         type = self.get_attr('type').content()
 113
 114         try:
 115             if type == 'dict':
 116                 return dict([ (item.get_attr('key'), item.as_pickled()) for item in self.get_list('item') ])
 117             elif type == 'list':
 118                 return [item.as_pickled() for item in self.get_list('item')]
 119             elif type == 'bool':
 120                 return self.content().lower() == 'true'
 121             elif type in RowElement.TYPES_MAP:
 122                 return RowElement.TYPES_MAP[type](self.content())
 123             else:
 124                 return self.content()
 125         except:
 126             return default
 127
 128
 129
 130
 131 class TableHandler(ContentHandler):
 132     def __init__(self, root_name, row_name, callback, callback_args = [], ping = None):
 133         self.root_name = root_name.lower()
 134         self.row_name = row_name.lower()
 135         self.callback = callback
 136         self.callback_args = callback_args
 137         self.ping = ping
 138
 139         self._reset()
 140
 141     def _reset(self):
 142         self.curr_element = None
 143         self.in_tag = None
 144
 145     def startElement(self, name, attrs):
 146         name = name.lower()
 147
 148         if name == self.root_name.lower():
 149             pass
 150         elif name == self.row_name:
 151             self.curr_element = RowElement(name, attrs)
 152         else:
 153             self.curr_element = RowElement(name, attrs, self.curr_element)
 154
 155     def characters(self, ch):
 156         if self.curr_element:
 157             self.curr_element.add_to_content(ch)
 158
 159     def endElement(self, name):
 160         name = name.lower()
 161
 162         if name == self.root_name:
 163             pass
 164         elif name == self.row_name:
 165             self.callback(self.curr_element, *self.callback_args)
 166             if self.ping:
 167                 self.ping()
 168
 169             self._reset()
 170         else:
 171             self.curr_element = self.curr_element.parent
 172
 173
 174 class SaxErrorHandler(ErrorHandler):
 175     def error(self, e):
 176         raise e
 177
 178     def fatalError(self, e):
 179         raise e
 180
 181     def warning(self, e):
 182         raise e
 183
 184 def disable_triggers():
 185     if db.backend_name == "postgres":
 186         db.start_transaction()
 187         db.execute_many(commands.PG_DISABLE_TRIGGERS)
 188         db.commit_transaction()
 189
 190 def enable_triggers():
 191     if db.backend_name == "postgres":
 192         db.start_transaction()
 193         db.execute_many(commands.PG_ENABLE_TRIGGERS)
 194         db.commit_transaction()
 195
 196 def reset_sequences():
 197     if db.backend_name == "postgres":
 198         db.start_transaction()
 199         db.execute_many(commands.PG_SEQUENCE_RESETS)
 200         db.commit_transaction()
 201
 202 FILE_HANDLERS = []
 203
 204 def start_import(fname, user):
 205
 206     start_time = datetime.datetime.now()
 207     steps = [s for s in FILE_HANDLERS]
 208
 209     with open(os.path.join(TMP_FOLDER, 'backup.inf'), 'r') as inffile:
 210         inf = ConfigParser.SafeConfigParser()
 211         inf.readfp(inffile)
 212
 213         state = dict([(s['id'], {
 214             'status': _('Queued'), 'count': int(inf.get(META_INF_SECTION, s['id'])), 'parsed': 0
 215         }) for s in steps] + [
 216             ('overall', {
 217                 'status': _('Starting'), 'count': int(inf.get(META_INF_SECTION, 'overall')), 'parsed': 0
 218             })
 219         ])
 220
 221     full_state = dict(running=True, state=state, time_started="")
 222
 223     def set_state():
 224         full_state['time_started'] = diff_date(start_time)
 225         cache.set(CACHE_KEY, full_state)
 226
 227     set_state()
 228
 229     def ping_state(name):
 230         state[name]['parsed'] += 1
 231         state['overall']['parsed'] += 1
 232         set_state()
 233
 234     def run(fn, name):
 235         def ping():
 236             ping_state(name)
 237
 238         state['overall']['status'] = _('Importing %s') % s['name']
 239         state[name]['status'] = _('Importing')
 240
 241
 242         fn(TMP_FOLDER, user, ping)
 243
 244         state[name]['status'] = _('Done')
 245
 246         set_state()
 247
 248         return fname
 249
 250     #dump = tarfile.open(fname, 'r')
 251     #dump.extractall(TMP_FOLDER)
 252
 253     try:
 254
 255         disable_triggers()
 256         db.start_transaction()
 257
 258         for h in FILE_HANDLERS:
 259             run(h['fn'], h['id'])
 260
 261         raise Exception
 262         db.commit_transaction()
 263         enable_triggers()
 264
 265         reset_sequences()
 266     except Exception, e:
 267         full_state['running'] = False
 268         full_state['errors'] = "%s: %s" % (e.__class__.__name__, unicode(e))
 269         set_state()
 270
 271         import traceback
 272         logging.error("Error executing xml import: \n %s" % (traceback.format_exc()))
 273
 274 def file_handler(file_name, root_tag, el_tag, name, args_handler=None, pre_callback=None, post_callback=None):
 275     def decorator(fn):
 276         def decorated(location, current_user, ping):
 277             if pre_callback:
 278                 pre_callback(current_user)
 279
 280             if (args_handler):
 281                 args = args_handler(current_user)
 282             else:
 283                 args = []
 284
 285             parser = make_parser()
 286             handler = TableHandler(root_tag, el_tag, fn, args, ping)
 287             parser.setContentHandler(handler)
 288             #parser.setErrorHandler(SaxErrorHandler())
 289
 290             parser.parse(os.path.join(location, file_name))
 291
 292             if post_callback:
 293                 post_callback()
 294
 295         FILE_HANDLERS.append(dict(id=root_tag, name=name, fn=decorated))
 296         return decorated
 297     return decorator
 298
 299
 300 @file_handler('users.xml', 'users', 'user', _('Users'), args_handler=lambda u: [u])
 301 def user_import(row, current_user):
 302     existent = False
 303
 304     if str(current_user.id) == row.getc('id'):
 305         existent = True
 306
 307     roles = row.get('roles').get_listc('role')
 308     valid_email = row.get('email').get_attr('validated').as_bool()
 309     badges = row.get('badges')
 310
 311     user = orm.User(
 312             id           = row.getc('id'),
 313             username     = row.getc('username'),
 314             password     = row.getc('password'),
 315             email        = row.getc('email'),
 316             email_isvalid= valid_email,
 317             is_superuser = 'superuser' in roles,
 318             is_staff     = 'moderator' in roles,
 319             is_active    = True,
 320             date_joined  = row.get('joindate').as_datetime(),
 321             about         = row.getc('bio'),
 322             date_of_birth = row.get('birthdate').as_date(None),
 323             website       = row.getc('website'),
 324             reputation    = row.get('reputation').as_int(),
 325             gold          = badges.get_attr('gold').as_int(),
 326             silver        = badges.get_attr('silver').as_int(),
 327             bronze        = badges.get_attr('bronze').as_int(),
 328             real_name     = row.getc('realname'),
 329             location      = row.getc('location'),
 330     )
 331
 332     user.save()
 333
 334     authKeys = row.get('authKeys')
 335
 336     for key in authKeys.get_list('key'):
 337         orm.AuthKeyUserAssociation(user=user, key=key.getc('key'), provider=key.getc('provider')).save()
 338
 339     notifications = row.get('notifications')
 340
 341     attributes = dict([(str(k), v.as_bool() and 'i' or 'n') for k, v in notifications.get('notify').attrs.items()])
 342     attributes.update(dict([(str(k), v.as_bool()) for k, v in notifications.get('autoSubscribe').attrs.items()]))
 343     attributes.update(dict([(str("notify_%s" % k), v.as_bool()) for k, v in notifications.get('notifyOnSubscribed').attrs.items()]))
 344
 345     ss = orm.SubscriptionSettings(user=user, enable_notifications=notifications.get_attr('enabled').as_bool(), **attributes)
 346
 347     if existent:
 348         ss.id = current_user.subscription_settings.id
 349
 350     ss.save()
 351
 352
 353 def pre_tag_import(user):
 354     tag_import.tag_mappings={}
 355
 356
 357 @file_handler('tags.xml', 'tags', 'tag', _('Tags'), pre_callback=pre_tag_import)
 358 def tag_import(row):
 359     tag = orm.Tag(name=row.getc('name'), used_count=row.get('used').as_int(), created_by_id=row.get('author').as_int())
 360     tag.save()
 361     tag_import.tag_mappings[tag.name] = tag
 362
 363
 364 def post_node_import():
 365     tag_import.tag_mappings = None
 366
 367 @file_handler('nodes.xml', 'nodes', 'node', _('Nodes'), args_handler=lambda u: [tag_import.tag_mappings], post_callback=post_node_import)
 368 def node_import(row, tags):
 369
 370     ntags = []
 371
 372     for t in row.get('tags').get_list('tag'):
 373         ntags.append(tags[t.content()])
 374
 375     last_act = row.get('lastactivity')
 376
 377     node = orm.Node(
 378             id            = row.getc('id'),
 379             node_type     = row.getc('type'),
 380             author_id     = row.get('author').as_int(),
 381             added_at      = row.get('date').as_datetime(),
 382             parent_id     = row.get('parent').as_int(None),
 383             abs_parent_id = row.get('absparent').as_int(None),
 384             score         = row.get('score').as_int(0),
 385
 386             last_activity_by_id = last_act.get('by').as_int(None),
 387             last_activity_at    = last_act.get('at').as_datetime(None),
 388
 389             title         = row.getc('title'),
 390             body          = row.getc('body'),
 391             tagnames      = " ".join([t.name for t in ntags]),
 392
 393             marked        = row.get('marked').as_bool(),
 394             extra_ref_id  = row.get('extraRef').as_int(None),
 395             extra_count   = row.get('extraCount').as_int(0),
 396             extra         = row.get('extraData').as_pickled()
 397     )
 398
 399     node.save()
 400     node.tags = ntags
 401
 402     revisions = row.get('revisions')
 403     active = revisions.get_attr('active').as_int()
 404
 405     for r in revisions.get_list('revision'):
 406         rev = orm.NodeRevision(
 407             author_id = r.getc('author'),
 408             body = r.getc('body'),
 409             node = node,
 410             revised_at = r.get('date').as_datetime(),
 411             revision = r.get('number').as_int(),
 412             summary = r.getc('summary'),
 413             tagnames = " ".join(r.getc('tags').split(',')),
 414             title = r.getc('title'),
 415         )
 416
 417         rev.save()
 418         if rev.revision == active:
 419             active = rev
 420
 421     node.active_revision = active
 422     node.save()
 423
 424 POST_ACTION = {}
 425
 426 def post_action(*types):
 427     def decorator(fn):
 428         for t in types:
 429             POST_ACTION[t] = fn
 430         return fn
 431     return decorator
 432
 433 def post_action_import_callback():
 434     with_state = orm.Node.objects.filter(id__in=orm.NodeState.objects.values_list('node_id', flat=True).distinct())
 435
 436     for n in with_state:
 437         n.state_string = "".join(["(%s)" % s for s in n.states.values_list('state_type')])
 438         n.save()
 439
 440 @file_handler('actions.xml', 'actions', 'action', _('Actions'), post_callback=post_action_import_callback)
 441 def actions_import(row):
 442     action = orm.Action(
 443         id           = row.get('id').as_int(),
 444         action_type  = row.getc('type'),
 445         action_date  = row.get('date').as_datetime(),
 446         node_id      = row.get('node').as_int(None),
 447         user_id      = row.get('user').as_int(),
 448         real_user_id = row.get('realUser').as_int(None),
 449         ip           = row.getc('ip'),
 450         extra        = row.get('extraData').as_pickled(),
 451     )
 452
 453     canceled = row.get('canceled')
 454     if canceled.get_attr('state').as_bool():
 455         action.canceled = True
 456         action.canceled_by_id = canceled.get('user').as_int()
 457         #action.canceled_at = canceled.get('date').as_datetime(),
 458         action.canceled_ip = canceled.getc('ip')
 459
 460     action.save()
 461
 462     for r in row.get('reputes').get_list('repute'):
 463         by_canceled = r.get_attr('byCanceled').as_bool()
 464
 465         orm.ActionRepute(
 466             action = action,
 467             user_id = r.get('user').as_int(),
 468             value = r.get('value').as_int(),
 469
 470             date = by_canceled and action.canceled_at or action.action_date,
 471             by_canceled = by_canceled
 472         ).save()
 473
 474     if (not action.canceled) and (action.action_type in POST_ACTION):
 475         POST_ACTION[action.action_type](row, action)
 476
 477
 478
 479
 480 @post_action('voteup', 'votedown', 'voteupcomment')
 481 def vote_action(row, action):
 482     orm.Vote(user_id=action.user_id, node_id=action.node_id, action=action,
 483              voted_at=action.action_date, value=(action.action_type != 'votedown') and 1 or -1).save()
 484
 485 def state_action(state):
 486     def fn(row, action):
 487         orm.NodeState(
 488             state_type = state,
 489             node_id = action.node_id,
 490             action = action
 491         ).save()
 492     return fn
 493
 494 post_action('wikify')(state_action('wiki'))
 495 post_action('delete')(state_action('deleted'))
 496 post_action('acceptanswer')(state_action('accepted'))
 497 post_action('publish')(state_action('published'))
 498
 499
 500 @post_action('flag')
 501 def flag_action(row, action):
 502     orm.Flag(user_id=action.user_id, node_id=action.node_id, action=action, reason=action.extra).save()
 503
 504
 505 def award_import_args(user):
 506     return [ dict([ (b.cls, b) for b in orm.Badge.objects.all() ]) ]
 507
 508
 509 @file_handler('awards.xml', 'awards', 'award', _('Awards'), args_handler=award_import_args)
 510 def awards_import(row, badges):
 511     try:
 512         badge_type = badges.get(row.getc('badge'), None)
 513
 514         if not badge_type:
 515             return
 516
 517         award = orm.Award(
 518             user_id = row.get('user').as_int(),
 519             badge = badges[row.getc('badge')],
 520             node_id = row.get('node').as_int(None),
 521             action_id = row.get('action').as_int(None),
 522             trigger_id = row.get('trigger').as_int(None)
 523         ).save()
 524     except Exception, e:
 525         return
 526
 527
 528 @file_handler('settings.xml', 'settings', 'setting', _('Settings'))
 529 def settings_import(row):
 530     orm.KeyValue(key=row.getc('key'), value=row.get('value').as_pickled())
 531
 532
 533
 534
 535
 536
 537
 538
 539