flask开发restful api系列(2)

　　继续上一章所讲，上一章我们最后面说道，虽然这个是很小的程序，但还有好几个要优化的地方。先复制一下老的view.py代码。

 # coding:utf-8
 from flask import Flask, request, jsonify
 from model import User, db_session
 import hashlib
 import time
 import redis

 app = Flask(__name__)
 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123')

 @app.route('/')
 def hello_world():
     return 'Hello World!'

 @app.route('/login', methods=['POST'])
 def login():
     phone_number = request.get_json().get('phone_number')
     password = request.get_json().get('password')
     user = User.query.filter_by(phone_number=phone_number).first()
     if not user:
         return jsonify({'code': 0, 'message': '没有此用户'})

     if user.password != password:
         return jsonify({'code': 0, 'message': '密码错误'})

     m = hashlib.md5()
     m.update(phone_number)
     m.update(password)
     m.update(str(int(time.time())))
     token = m.hexdigest()

     redis_store.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1})
     redis_store.set('token:%s' % token, user.phone_number)
     redis_store.expire('token:%s' % token, 3600*24*30)

     return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token})

 @app.route('/user')
 def user():
     token = request.headers.get('token')
     if not token:
         return jsonify({'code': 0, 'message': '需要验证'})
     phone_number = redis_store.get('token:%s' % token)
     if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'):
         return jsonify({'code': 2, 'message': '验证信息错误'})

     nickname = redis_store.hget('user:%s' % phone_number, 'nickname')
     return jsonify({'code': 1, 'nickname': nickname, 'phone_number': phone_number})

 @app.route('/logout')
 def logout():
     token = request.headers.get('token')
     if not token:
         return jsonify({'code': 0, 'message': '需要验证'})
     phone_number = redis_store.get('token:%s' % token)
     if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'):
         return jsonify({'code': 2, 'message': '验证信息错误'})

     redis_store.delete('token:%s' % token)
     redis_store.hmset('user:%s' % phone_number, {'app_online': 0})
     return jsonify({'code': 1, 'message': '成功注销'})

 @app.teardown_request
 def handle_teardown_request(exception):
     db_session.remove()

 if __name__ == '__main__':
     app.run(debug=True, host='0.0.0.0', port=5001)

其中验证token的方法，已经重叠了，python教我们，永远不要重复自己的代码，这是很丑陋的行为。今天我们把它换成一个装饰器，然后再把redis调整一下，看看代码会不会简洁很多。

 # coding:utf-8
 from flask import Flask, request, jsonify
 from model import User, db_session
 import hashlib
 import time
 import redis
 from functools import wraps

 app = Flask(__name__)
 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123')

 def login_check(f):
     @wraps(f)
     def decorator(*args, **kwargs):
         token = request.headers.get('token')
         if not token:
             return jsonify({'code': 0, 'message': '需要验证'})

         phone_number = redis_store.get('token:%s' % token)
         if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'):
             return jsonify({'code': 2, 'message': '验证信息错误'})

         return f(*args, **kwargs)
     return decorator

 @app.route('/login', methods=['POST'])
 def login():
     phone_number = request.get_json().get('phone_number')
     password = request.get_json().get('password')
     user = User.query.filter_by(phone_number=phone_number).first()
     if not user:
         return jsonify({'code': 0, 'message': '没有此用户'})

     if user.password != password:
         return jsonify({'code': 0, 'message': '密码错误'})

     m = hashlib.md5()
     m.update(phone_number)
     m.update(password)
     m.update(str(int(time.time())))
     token = m.hexdigest()

     pipeline = redis_store.pipeline()
     pipeline.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1})
     pipeline.set('token:%s' % token, user.phone_number)
     pipeline.expire('token:%s' % token, 3600*24*30)
     pipeline.execute()

     return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token})

 @app.route('/user')
 @login_check
 def user():
     token = request.headers.get('token')
     phone_number = redis_store.get('token:%s' % token)

     nickname = redis_store.hget('user:%s' % phone_number, 'nickname')
     return jsonify({'code': 1, 'nickname': nickname, 'phone_number': phone_number})

 @app.route('/logout')
 @login_check
 def logout():
     token = request.headers.get('token')
     phone_number = redis_store.get('token:%s' % token)

     pipeline = redis_store.pipeline()
     pipeline.delete('token:%s' % token)
     pipeline.hmset('user:%s' % phone_number, {'app_online': 0})
     pipeline.execute()
     return jsonify({'code': 1, 'message': '成功注销'})

 @app.teardown_request
 def handle_teardown_request(exception):
     db_session.remove()

 if __name__ == '__main__':
     app.run(debug=True, host='0.0.0.0', port=5001)

加了一个装饰器，是不是简洁了很多？每次需要验证的时候，只需要一个login_check就可以了，这样就变得非常简洁，而且脉络清晰。redis也改成了管道执行，pipeline，防止执行到一半，被掐断。

可是，可是，我还是觉得不简洁，看user, logout的代码中重复的地方。

token = request.headers.get('token')
phone_number = redis_store.get('token:%s' % token)

每次都有这两句，要是将来还有其他值怎么办？上面不刚说，永远不要重复自己的代码吗？

好，我们再写一个函数，看下面代码

 # coding:utf-8
 from flask import Flask, request, jsonify, g
 from model import User, db_session
 import hashlib
 import time
 import redis
 from functools import wraps

 app = Flask(__name__)
 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123')

 def login_check(f):
     @wraps(f)
     def decorator(*args, **kwargs):
         token = request.headers.get('token')
         if not token:
             return jsonify({'code': 0, 'message': '需要验证'})

         phone_number = redis_store.get('token:%s' % token)
         if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'):
             return jsonify({'code': 2, 'message': '验证信息错误'})

         return f(*args, **kwargs)
     return decorator

 @app.before_request
 def before_request():
     token = request.headers.get('token')
     phone_number = redis_store.get('token:%s' % token)
     if phone_number:
         g.current_user = User.query.filter_by(phone_number=phone_number).first()
         g.token = token
     return

 @app.route('/login', methods=['POST'])
 def login():
     phone_number = request.get_json().get('phone_number')
     password = request.get_json().get('password')
     user = User.query.filter_by(phone_number=phone_number).first()
     if not user:
         return jsonify({'code': 0, 'message': '没有此用户'})

     if user.password != password:
         return jsonify({'code': 0, 'message': '密码错误'})

     m = hashlib.md5()
     m.update(phone_number)
     m.update(password)
     m.update(str(int(time.time())))
     token = m.hexdigest()

     pipeline = redis_store.pipeline()
     pipeline.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1})
     pipeline.set('token:%s' % token, user.phone_number)
     pipeline.expire('token:%s' % token, 3600*24*30)
     pipeline.execute()

     return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token})

 @app.route('/user')
 @login_check
 def user():
     user = g.current_user

     nickname = redis_store.hget('user:%s' % user.phone_number, 'nickname')
     return jsonify({'code': 1, 'nickname': nickname, 'phone_number': user.phone_number})

 @app.route('/logout')
 @login_check
 def logout():
     user = g.current_user

     pipeline = redis_store.pipeline()
     pipeline.delete('token:%s' % g.token)
     pipeline.hmset('user:%s' % user.phone_number, {'app_online': 0})
     pipeline.execute()
     return jsonify({'code': 1, 'message': '成功注销'})

 @app.teardown_request
 def handle_teardown_request(exception):
     db_session.remove()

 if __name__ == '__main__':
     app.run(debug=True, host='0.0.0.0', port=5001)

我们在代码中加了一个before_request，这个函数就是在每个request发起的时候，如果已经验证了，我们把当前g.current_user和g.token设置一下，这样每次需要获取当前用户的时候，直接找g.current_user就可以了，是不是简单了太多太多？好了，今天到此为止，下一章，我们讲怎么利用alembic修改数据库。