继续上一章所讲，上一章我们最后面说道，虽然这个是很小的程序，但还有好几个要优化的地方。先复制一下老的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修改数据库。