flask开发restful api系列(2)

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

 1 # coding:utf-8
 2 from flask import Flask, request, jsonify
 3 from model import User, db_session
 4 import hashlib
 5 import time
 6 import redis
 7 
 8 app = Flask(__name__)
 9 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123')
10 
11 
12 @app.route('/')
13 def hello_world():
14     return 'Hello World!'
15 
16 
17 @app.route('/login', methods=['POST'])
18 def login():
19     phone_number = request.get_json().get('phone_number')
20     password = request.get_json().get('password')
21     user = User.query.filter_by(phone_number=phone_number).first()
22     if not user:
23         return jsonify({'code': 0, 'message': '没有此用户'})
24 
25     if user.password != password:
26         return jsonify({'code': 0, 'message': '密码错误'})
27 
28     m = hashlib.md5()
29     m.update(phone_number)
30     m.update(password)
31     m.update(str(int(time.time())))
32     token = m.hexdigest()
33 
34     redis_store.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1})
35     redis_store.set('token:%s' % token, user.phone_number)
36     redis_store.expire('token:%s' % token, 3600*24*30)
37 
38     return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token})
39 
40 
41 @app.route('/user')
42 def user():
43     token = request.headers.get('token')
44     if not token:
45         return jsonify({'code': 0, 'message': '需要验证'})
46     phone_number = redis_store.get('token:%s' % token)
47     if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'):
48         return jsonify({'code': 2, 'message': '验证信息错误'})
49 
50     nickname = redis_store.hget('user:%s' % phone_number, 'nickname')
51     return jsonify({'code': 1, 'nickname': nickname, 'phone_number': phone_number})
52 
53 
54 @app.route('/logout')
55 def logout():
56     token = request.headers.get('token')
57     if not token:
58         return jsonify({'code': 0, 'message': '需要验证'})
59     phone_number = redis_store.get('token:%s' % token)
60     if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'):
61         return jsonify({'code': 2, 'message': '验证信息错误'})
62 
63     redis_store.delete('token:%s' % token)
64     redis_store.hmset('user:%s' % phone_number, {'app_online': 0})
65     return jsonify({'code': 1, 'message': '成功注销'})
66 
67 
68 @app.teardown_request
69 def handle_teardown_request(exception):
70     db_session.remove()
71 
72 if __name__ == '__main__':
73     app.run(debug=True, host='0.0.0.0', port=5001)

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

 1 # coding:utf-8
 2 from flask import Flask, request, jsonify
 3 from model import User, db_session
 4 import hashlib
 5 import time
 6 import redis
 7 from functools import wraps
 8 
 9 app = Flask(__name__)
10 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123')
11 
12 
13 def login_check(f):
14     @wraps(f)
15     def decorator(*args, **kwargs):
16         token = request.headers.get('token')
17         if not token:
18             return jsonify({'code': 0, 'message': '需要验证'})
19         
20         phone_number = redis_store.get('token:%s' % token)
21         if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'):
22             return jsonify({'code': 2, 'message': '验证信息错误'})
23 
24         return f(*args, **kwargs)
25     return decorator
26 
27 
28 @app.route('/login', methods=['POST'])
29 def login():
30     phone_number = request.get_json().get('phone_number')
31     password = request.get_json().get('password')
32     user = User.query.filter_by(phone_number=phone_number).first()
33     if not user:
34         return jsonify({'code': 0, 'message': '没有此用户'})
35 
36     if user.password != password:
37         return jsonify({'code': 0, 'message': '密码错误'})
38 
39     m = hashlib.md5()
40     m.update(phone_number)
41     m.update(password)
42     m.update(str(int(time.time())))
43     token = m.hexdigest()
44     
45     pipeline = redis_store.pipeline()
46     pipeline.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1})
47     pipeline.set('token:%s' % token, user.phone_number)
48     pipeline.expire('token:%s' % token, 3600*24*30)
49     pipeline.execute()
50 
51     return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token})
52 
53 
54 @app.route('/user')
55 @login_check
56 def user():
57     token = request.headers.get('token')
58     phone_number = redis_store.get('token:%s' % token)
59 
60     nickname = redis_store.hget('user:%s' % phone_number, 'nickname')
61     return jsonify({'code': 1, 'nickname': nickname, 'phone_number': phone_number})
62 
63 
64 @app.route('/logout')
65 @login_check
66 def logout():
67     token = request.headers.get('token')
68     phone_number = redis_store.get('token:%s' % token)
69     
70     pipeline = redis_store.pipeline()
71     pipeline.delete('token:%s' % token)
72     pipeline.hmset('user:%s' % phone_number, {'app_online': 0})
73     pipeline.execute()
74     return jsonify({'code': 1, 'message': '成功注销'})
75 
76 
77 @app.teardown_request
78 def handle_teardown_request(exception):
79     db_session.remove()
80 
81 if __name__ == '__main__':
82     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)

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

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

 1 # coding:utf-8
 2 from flask import Flask, request, jsonify, g
 3 from model import User, db_session
 4 import hashlib
 5 import time
 6 import redis
 7 from functools import wraps
 8 
 9 app = Flask(__name__)
10 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123')
11 
12 
13 def login_check(f):
14     @wraps(f)
15     def decorator(*args, **kwargs):
16         token = request.headers.get('token')
17         if not token:
18             return jsonify({'code': 0, 'message': '需要验证'})
19         
20         phone_number = redis_store.get('token:%s' % token)
21         if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'):
22             return jsonify({'code': 2, 'message': '验证信息错误'})
23 
24         return f(*args, **kwargs)
25     return decorator
26 
27 
28 @app.before_request
29 def before_request():
30     token = request.headers.get('token')
31     phone_number = redis_store.get('token:%s' % token)
32     if phone_number:
33         g.current_user = User.query.filter_by(phone_number=phone_number).first()
34         g.token = token
35     return
36 
37 
38 @app.route('/login', methods=['POST'])
39 def login():
40     phone_number = request.get_json().get('phone_number')
41     password = request.get_json().get('password')
42     user = User.query.filter_by(phone_number=phone_number).first()
43     if not user:
44         return jsonify({'code': 0, 'message': '没有此用户'})
45 
46     if user.password != password:
47         return jsonify({'code': 0, 'message': '密码错误'})
48 
49     m = hashlib.md5()
50     m.update(phone_number)
51     m.update(password)
52     m.update(str(int(time.time())))
53     token = m.hexdigest()
54 
55     pipeline = redis_store.pipeline()
56     pipeline.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1})
57     pipeline.set('token:%s' % token, user.phone_number)
58     pipeline.expire('token:%s' % token, 3600*24*30)
59     pipeline.execute()
60 
61     return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token})
62 
63 
64 @app.route('/user')
65 @login_check
66 def user():
67     user = g.current_user
68 
69     nickname = redis_store.hget('user:%s' % user.phone_number, 'nickname')
70     return jsonify({'code': 1, 'nickname': nickname, 'phone_number': user.phone_number})
71 
72 
73 @app.route('/logout')
74 @login_check
75 def logout():
76     user = g.current_user
77 
78     pipeline = redis_store.pipeline()
79     pipeline.delete('token:%s' % g.token)
80     pipeline.hmset('user:%s' % user.phone_number, {'app_online': 0})
81     pipeline.execute()
82     return jsonify({'code': 1, 'message': '成功注销'})
83 
84 
85 @app.teardown_request
86 def handle_teardown_request(exception):
87     db_session.remove()
88 
89 if __name__ == '__main__':
90     app.run(debug=True, host='0.0.0.0', port=5001)

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

posted @ 2016-03-28 18:32  月儿弯弯0204  阅读(3104)  评论(0编辑  收藏  举报