update_db.py

#!/usr/bin/python
# -*- coding: utf8 -*-

# Author: Hu Ziming <[email protected]>
#
# License:
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#   See the GNU General Public License for more details.
#
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation; either version
#   2 of the License, or (at your option) any later version.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
#   02111-1307 USA


import yaml
import sys
import datetime

import sqlalchemy
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy import Table, Column, Integer, Unicode, UnicodeText, DateTime, MetaData, ForeignKey, PickleType
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relation, backref


Base = declarative_base()
class History (Base):
    __tablename__ = 'history'

    id           = Column(Integer, primary_key = True)
    title        = Column(Unicode)
    author       = Column(Unicode)
    content      = Column(UnicodeText)      # Source of content
    content_type = Column(UnicodeText)      # Type of content source (YAML, reST, html, txt, etc.)
    date         = Column(DateTime)
    type         = Column(Unicode)
    backup       = Column(PickleType)

    def __init__ (self, **kwd):
        self.title        = unicode(kwd['title'])
        self.author       = unicode(kwd['author'])
        self.content      = unicode(kwd['content'])
        self.content_type = unicode(kwd['content_type'])
        #self.date         = datetime.datetime(*map(int, kwd['date'].split('-')))
        self.date         = kwd['date']
        self.type         = unicode(kwd['type'])

        try:
            self.backup  = kwd['backup']
        except KeyError:
            self.backup = None

    def __repr__ (self):
        return "%s by %s" % (self.title, self.author)



def main():
    """ Read the content from YAML file and put it to the database"""

    # Create the session and connect the database
    #engine = create_engine('sqlite:///:memory:', echo = True)
    engine = create_engine('sqlite:///./foo.db')
    Session = sessionmaker (bind = engine)
    session = Session()

    # Get the content from YAML file
    fp = open(sys.argv[1])
    data = yaml.load_all(fp.read())
    fp.close()

    # Create the tables, need a judgement
    metadata = Base.metadata
    metadata.create_all(engine)

    # Add the data to the db file
    for rcd in data:
        try:
            if len(session.query(History.content).filter(History.content == rcd['content']).all()) > 0:
                print "Alread insterd, continue to the next record"
            else:
                session.add(History(**rcd))
        except KeyError:
            # continue to next record if error occurs
            print "Error Processing record:"
            print rcd
            continue

    session.commit()

if __name__ == '__main__':
    main()