Python models.DBSession类代码示例

def _new_user(args):
    """Handle adding a new user to the system.

    If you don't include the required info, it will prompt you for it

    """
    if not args.username:
        args.username = raw_input('username? ')

    if not args.email:
        args.email = raw_input('email address? ')

    if not args.username or not args.email:
        raise Exception('Must supply a username and email address')

    import transaction
    _init_sql(args)
    from bookie.models import DBSession
    sess = DBSession()

    u = User()
    u.username = unicode(args.username)
    passwd = get_random_word(8)
    u.password = passwd
    u.email = unicode(args.email)
    u.activated = True
    u.is_admin = False
    u.api_key = User.gen_api_key()

    print dict(u)
    print passwd

    sess.add(u)
    sess.flush()
    transaction.commit()

def new_user(request):
    """Add a new user to the system manually."""
    rdict = request.params

    u = User()

    u.username = unicode(rdict.get('username'))
    u.email = unicode(rdict.get('email'))
    passwd = get_random_word(8)
    u.password = passwd
    u.activated = True
    u.is_admin = False
    u.api_key = User.gen_api_key()

    try:
        DBSession.add(u)
        DBSession.flush()
        # We need to return the password since the admin added the user
        # manually.  This is only time we should have/give the original
        # password.
        ret = dict(u)
        ret['random_pass'] = passwd
        return _api_response(request, ret)

    except IntegrityError, exc:
        # We might try to add a user that already exists.
        LOG.error(exc)
        request.response.status_int = 400
        return _api_response(request, {
            'error': 'Bad Request: User exists.',
        })

def admin_bmark_remove(request):
    """Remove this bookmark from the system"""
    rdict = request.matchdict
    username = rdict.get('username')
    hash_id = rdict.get('hash_id')

    try:
        bmark = BmarkMgr.get_by_hash(hash_id,
                                     username=username)
        print bmark
        if bmark:
            DBSession.delete(bmark)
            return _api_response(request, {
                'message': "done",
            })
        else:
            return _api_response(request, {
                'error': 'Bookmark not found.',
            })

    except NoResultFound:
        request.response.status_code = 404
        return _api_response(request, {
            'error': 'Bookmark with hash id {0} not found.'.format(
                rdict['hash_id'])
        })

    def process(self):
        """Given a file, process it"""
        soup = BeautifulSoup(self.file_handle)
        count = 0

        ids = []
        for tag in soup.findAll("dt"):
            if "javascript:" in str(tag):
                continue

            # if we have a dd as next sibling, get it's content
            if tag.nextSibling and tag.nextSibling.name == "dd":
                extended = tag.nextSibling.text
            else:
                extended = u""

            link = tag.a

            # Skip any bookmarks with an attribute of PRIVATE.
            if link.has_key("PRIVATE"):
                continue

            import_add_date = float(link["add_date"])

            if import_add_date > 9999999999:
                # Remove microseconds from the timestamp
                import_add_date = import_add_date / 1000
            add_date = datetime.fromtimestamp(import_add_date)

            try:
                bmark = self.save_bookmark(
                    unicode(link["href"]),
                    unicode(link.text),
                    unicode(extended),
                    u" ".join(unicode(link.get("tags", "")).split(u",")),
                    dt=add_date,
                )
                count = count + 1
                DBSession.flush()
            except InvalidBookmark:
                bmark = None

            if bmark:
                ids.append(bmark.bid)

            if count % COMMIT_SIZE == 0:
                transaction.commit()

        # Commit any that are left since the last commit performed.
        transaction.commit()

        from bookie.bcelery import tasks

        # For each bookmark in this set that we saved, sign up to
        # fetch its content.
        for bid in ids:
            tasks.fetch_bmark_content.delay(bid)

        # Start a new transaction for the next grouping.
        transaction.begin()

    def test_update_post(self):
        """Updates allowed over the last one

        If you /post/add to an existing bookmark, the new data overrides the
        old data

        """
        self._get_good_request()

        # now build a new version of the request we can check
        session = DBSession()
        prms = {
                'url': u'http://google.com',
                'description': u'This is my updated google desc',
                'extended': 'updated extended notes about it in full form',
                'tags': u'python search updated',
                'api_key': API_KEY
        }

        req_params = urllib.urlencode(prms)
        self.testapp.get('/admin/delapi/posts/add?' + req_params)
        session.flush()

        res = Bmark.query.filter(Bmark.hash_id == GOOGLE_HASH).one()

        ok_('updated' in res.description,
                'Updated description took: ' + res.description)
        ok_('updated' in res.extended,
                'Updated extended took: ' + res.extended)
        ok_('python' in res.tags, 'Found the python tag in the bmark')
        ok_('search' in res.tags, 'Found the search tag in the bmark')
        ok_('updated' in res.tags, 'Found the updated tag in the bmark')

    def test_google_import(self):
        """Test that we can upload our google file"""
        self.testapp.post('/login',
                            params={
                                "login": "admin",
                                "password": "admin",
                                "form.submitted": "Log In",
                            },
                            status=302)

        session = DBSession()
        loc = os.path.dirname(__file__)
        del_file = open(os.path.join(loc, 'googlebookmarks.html'))
        res = DBSession.execute("SELECT api_key FROM users WHERE username = 'admin'").fetchone()
        API_KEY = res['api_key']

        post = {
            'api_key': API_KEY,
        }


        res = self.testapp.post('/admin/import',
                                params=post,
                                upload_files=[('import_file',
                                               'googlebookmarks.html',
                                               del_file.read())],
        )

        eq_(res.status, "302 Found",
            msg='Import status is 302 redirect by home, ' + res.status)

        session.flush()

        # test all the data we want to test after the import
        _google_data_test()

def db_init_bookmark():
    """install the initial bookmark in a new install"""
    require('hosts', provided_by=[sample])
    require('ini', provided_by=[sample])

    parse_ini(env["ini_file"])
    from datetime import datetime
    import transaction
    from bookie.models import initialize_sql
    from sqlalchemy import create_engine

    engine = create_engine(env.ini.get('app:bookie', 'sqlalchemy.url'))
    initialize_sql(engine)

    from bookie.models import DBSession
    from bookie.models import Bmark

    bmark_us = Bmark(u'http://bmark.us',
                     u'admin',
                     desc=u"Bookie Website",
                     ext= u"Bookie Documentation Home",
                     tags = u"bookmarks")

    bmark_us.stored = datetime.now()
    bmark_us.updated = datetime.now()
    DBSession.add(bmark_us)
    DBSession.flush()
    transaction.commit()

    def test_tag_with_space(self):
        """Test that we strip out spaces from tags and don't get empty tags

        """
        self._get_good_request()

        # now build a new version of the request we can check
        session = DBSession()
        prms = {
                'url': u'http://google.com',
                'description': u'This is my updated google desc',
                'extended': 'updated extended notes about it in full form',
                'tags': u'python  search updated ',
                'api_key': API_KEY
        }

        req_params = urllib.urlencode(prms)
        self.testapp.get('/admin/delapi/posts/add?' + req_params)
        session.flush()

        res = Bmark.query.filter(Bmark.hash_id == GOOGLE_HASH).one()

        ok_(len(res.tags) == 3,
                'Should only have 3 tags: ' + str([str(t) for t in res.tags]))

        for tag in res.tags:
            ok_(tag[0] != " ", "Tag should not start with a space")
            ok_(tag[-1] != " ", "Tag should not end with a space")

def new_user(username, email):
    """Add new user function, pass username, email

    :param username: string of new user
    :param email: string of new email

    """
    require('hosts', provided_by=[sample])
    require('ini', provided_by=[sample])

    parse_ini(env["ini_file"])

    import transaction
    from bookie.models import initialize_sql
    initialize_sql(dict(env.ini.items('app:main')))

    from bookie.models import DBSession
    from bookie.models.auth import get_random_word, User
    sess = DBSession()

    u = User()
    u.username = unicode(username)
    passwd = get_random_word(8)
    u.password = passwd
    u.email = unicode(email)
    u.activated = True
    u.is_admin = False
    u.api_key = User.gen_api_key()

    print dict(u)
    print passwd

    sess.add(u)
    sess.flush()
    transaction.commit()

    def process(self):
        """Given a file, process it"""
        if self.file_handle.closed:
            self.file_handle = open(self.file_handle.name)

        parsed = etree.parse(self.file_handle)
        count = 0

        ids = []
        for post in parsed.findall('post'):
            if 'javascript:' in post.get('href'):
                continue

            add_date = dateparser.parse(post.get('time'))

            try:
                bmark = self.save_bookmark(
                    post.get('href'),
                    post.get('description'),
                    post.get('extended'),
                    post.get('tag'),
                    dt=add_date)
                count = count + 1
                if bmark:
                    bmark.stored = bmark.stored.replace(tzinfo=None)
                    DBSession.flush()
            except InvalidBookmark, exc:
                bmark = None

            if bmark:
                ids.append(bmark.bid)

            if count % COMMIT_SIZE == 0:
                transaction.commit()

def upgrade(migrate_engine):
    """By default start all installs out with a bookmark to bmark.us

    It's cheesy, but helps with testing to have some base data and is good for
    pubbing the links

    """
    from datetime import datetime
    import transaction
    from bookie.models import initialize_sql

    initialize_sql(migrate_engine)

    from bookie.models import DBSession
    from bookie.models import Bmark

    bmark_us = Bmark('http://bmark.us',
                     desc="Bookie Website",
                     ext= "Bookie Documentation Home",
                     tags = "bookmarks")

    bmark_us.stored = datetime.now()
    bmark_us.updated = datetime.now()
    DBSession.add(bmark_us)
    DBSession.flush()
    transaction.commit()

 def tearDown(self):
     """Regular tear down method"""
     session = DBSession()
     Bmark.query.delete()
     Tag.query.delete()
     session.execute(bmarks_tags.delete())
     session.flush()
     transaction.commit()

    def import_bmarks(self):
        """Allow users to upload a bookmark export file for processing"""
        username = self.matchdict.get('username')

        # if auth fails, it'll raise an HTTPForbidden exception
        with ReqAuthorize(self.request):
            data = {}
            post = self.POST

            # We can't let them submit multiple times, check if this user has
            # an import in process.
            if ImportQueueMgr.get(username=username, status=NEW):
                # They have an import, get the information about it and shoot
                # to the template.
                return {
                    'existing': True,
                    'import_stats': ImportQueueMgr.get_details(
                        username=username)
                }

            if post:
                # we have some posted values
                files = post.get('import_file', None)

                if files is not None:
                    storage_dir_tpl = self.settings.get('import_files',
                                                        '/tmp/bookie')
                    storage_dir = storage_dir_tpl.format(
                        here=self.settings.get('app_root'))

                    out_fname = store_import_file(storage_dir, username, files)

                    # Mark the system that there's a pending import that needs
                    # to be completed
                    q = ImportQueue(username, unicode(out_fname))
                    DBSession.add(q)
                    DBSession.flush()
                    # Schedule a task to start this import job.
                    tasks.importer_process.delay(q.id)

                    return HTTPFound(
                        location=self.request.route_url('user_import',
                                                        username=username))
                else:
                    msg = self.request.session.pop_flash()
                    if msg:
                        data['error'] = msg
                    else:
                        data['error'] = None

                return data
            else:
                # we need to see if they've got
                # just display the form
                return {
                    'existing': False
                }

 def tearDown(self):
     """Tear down each test"""
     testing.tearDown()
     session = DBSession()
     Bmark.query.delete()
     Tag.query.delete()
     session.execute(bmarks_tags.delete())
     session.flush()
     transaction.commit()

    def test_total_ct(self):
        """Verify that our total count method is working"""
        ct = 5
        for i in range(ct):
            t = Tag(gen_random_word(10))
            DBSession.add(t)

        ct = TagMgr.count()
        eq_(5, ct, 'We should have a total of 5: ' + str(ct))