added request decorators and function annotations

spotify/api/api.py

@@ -4,6 +4,7 @@ import os
 import datetime
 import json
 import logging
+import functools
 
 from google.cloud import firestore
 from google.cloud import tasks_v2
@@ -24,10 +25,64 @@ task_path = tasker.queue_path('sarsooxyz', 'europe-west2', 'spotify-executions')
 logger = logging.getLogger(__name__)
 
 
-@blueprint.route('/playlists', methods=['GET'])
+def login_required(func):
-def get_playlists():
+    @functools.wraps(func)
-
+    def wrapper(*args, **kwargs):
         if 'username' in session:
+            return func(*args, **kwargs)
+        else:
+            logger.warning('user not logged in')
+            return jsonify({'error': 'not logged in'}), 401
+    return wrapper
+
+
+def admin_required(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        user_dict = database.get_user_doc_ref(session['username']).get().to_dict()
+
+        if user_dict:
+            if user_dict['type'] == 'admin':
+                return func(*args, **kwargs)
+            else:
+                logger.warning(f'{user_dict["username"]} not authorized')
+                return jsonify({'status': 'error', 'message': 'unauthorized'}), 401
+        else:
+            logger.warning('user not logged in')
+            return jsonify({'error': 'not logged in'}), 401
+
+    return wrapper
+
+
+def gae_cron(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+
+        if request.headers.get('X-Appengine-Cron', None):
+            return func(*args, **kwargs)
+        else:
+            logger.warning('user not logged in')
+            return jsonify({'status': 'error', 'message': 'unauthorised'}), 401
+
+    return wrapper
+
+
+def cloud_task(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+
+        if request.headers.get('X-AppEngine-QueueName', None):
+            return func(*args, **kwargs)
+        else:
+            logger.warning('non tasks request')
+            return jsonify({'status': 'error', 'message': 'unauthorised'}), 401
+
+    return wrapper
+
+
+@blueprint.route('/playlists', methods=['GET'])
+@login_required
+def get_playlists():
 
     pulled_user = database.get_user_doc_ref(session['username'])
 
@@ -45,16 +100,11 @@ def get_playlists():
 
     return jsonify(response), 200
 
-    else:
-        logger.warning('user not logged in')
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/playlist', methods=['GET', 'POST', 'PUT', 'DELETE'])
+@login_required
 def playlist():
 
-    if 'username' in session:
-
     user_ref = database.get_user_doc_ref(session['username'])
     playlists = user_ref.collection(u'playlists')
 
@@ -219,16 +269,11 @@ def playlist():
 
             return jsonify({"message": 'playlist updated', "status": "success"}), 200
 
-    else:
-        logger.warning('user not logged in')
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/user', methods=['GET', 'POST'])
+@login_required
 def user():
 
-    if 'username' in session:
-
     if request.method == 'GET':
 
         pulled_user = database.get_user_doc_ref(session['username']).get().to_dict()
@@ -281,19 +326,12 @@ def user():
 
         return jsonify({'message': 'account updated', 'status': 'succeeded'}), 200
 
-    else:
-        logger.warning('user not logged in')
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/users', methods=['GET'])
+@login_required
+@admin_required
 def users():
 
-    if 'username' in session:
-
-        if database.get_user_doc_ref(session['username']).get().to_dict()['type'] != 'admin':
-            return jsonify({'status': 'unauthorised'}), 401
-
     dic = {
         'accounts': []
     }
@@ -312,17 +350,13 @@ def users():
 
     return jsonify(dic), 200
 
-    else:
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/user/password', methods=['POST'])
+@login_required
 def change_password():
 
     request_json = request.get_json()
 
-    if 'username' in session:
-
     if 'new_password' in request_json and 'current_password' in request_json:
 
         if len(request_json['new_password']) == 0:
@@ -347,15 +381,11 @@ def change_password():
     else:
         return jsonify({'error': 'malformed request, no old_password/new_password'}), 400
 
-    else:
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/playlist/play', methods=['POST'])
+@login_required
 def play_playlist():
 
-    if 'username' in session:
-
     request_json = request.get_json()
 
     request_parts = request_json.get('parts', None)
@@ -406,14 +436,10 @@ def play_playlist():
         logger.error(f'no playlists/parts {session["username"]}')
         return jsonify({'error': 'insufficient playlist sources'}), 400
 
-    else:
-        logger.warning('user not logged in')
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/playlist/play/task', methods=['POST'])
+@cloud_task
 def play_playlist_task():
-    if request.headers.get('X-AppEngine-QueueName', None):
     payload = request.get_data(as_text=True)
     if payload:
         payload = json.loads(payload)
@@ -431,16 +457,12 @@ def play_playlist_task():
                            add_last_month=payload['add_last_month'])
 
         return jsonify({'message': 'executed playlist', 'status': 'success'}), 200
-    else:
-        logger.warning('non tasks request')
-        return jsonify({'error': 'unauthorized'}), 401
 
 
 @blueprint.route('/playlist/run', methods=['GET'])
+@login_required
 def run_playlist():
 
-    if 'username' in session:
-
     playlist_name = request.args.get('name', None)
 
     if playlist_name:
@@ -456,15 +478,11 @@ def run_playlist():
         logger.warning('no playlist requested')
         return jsonify({"error": 'no name requested'}), 400
 
-    else:
-        logger.warning('user not logged in')
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/playlist/run/task', methods=['POST'])
+@cloud_task
 def run_playlist_task():
 
-    if request.headers.get('X-AppEngine-QueueName', None):
     payload = request.get_data(as_text=True)
     if payload:
         payload = json.loads(payload)
@@ -474,16 +492,12 @@ def run_playlist_task():
         run_user_playlist(payload['username'], payload['name'])
 
         return jsonify({'message': 'executed playlist', 'status': 'success'}), 200
-    else:
-        logger.warning('non tasks request')
-        return jsonify({'error': 'unauthorized'}), 401
 
 
 @blueprint.route('/playlist/run/user', methods=['GET'])
+@login_required
 def run_user():
 
-    if 'username' in session:
-
     if database.get_user_doc_ref(session['username']).get().to_dict()['type'] == 'admin':
         user_name = request.args.get('username', session['username'])
     else:
@@ -493,50 +507,32 @@ def run_user():
 
     return jsonify({'message': 'executed user', 'status': 'success'}), 200
 
-    else:
-        logger.warning('user not logged in')
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/playlist/run/user/task', methods=['POST'])
+@cloud_task
 def run_user_task():
 
-    if request.headers.get('X-AppEngine-QueueName', None):
     payload = request.get_data(as_text=True)
     if payload:
         execute_user(payload)
         return jsonify({'message': 'executed user', 'status': 'success'}), 200
-    else:
-        logger.warning('non tasks request')
-        return jsonify({'error': 'unauthorized'}), 401
 
 
 @blueprint.route('/playlist/run/users', methods=['GET'])
+@login_required
+@admin_required
 def run_users():
 
-    if 'username' in session:
-
-        if database.get_user_doc_ref(session['username']).get().to_dict()['type'] != 'admin':
-            return jsonify({'status': 'error', 'message': 'unauthorized'}), 401
-
     execute_all_users()
-
     return jsonify({'message': 'executed all users', 'status': 'success'}), 200
 
-    else:
-        logger.warning('user not logged in')
-        return jsonify({'error': 'not logged in'}), 401
-
 
 @blueprint.route('/playlist/run/users/cron', methods=['GET'])
+@gae_cron
 def run_users_cron():
 
-    if request.headers.get('X-Appengine-Cron', None):
     execute_all_users()
     return jsonify({'status': 'success'}), 200
-    else:
-        logger.warning('user not logged in')
-        return jsonify({'status': 'error', 'message': 'unauthorised'}), 401
 
 
 def execute_all_users():

spotify/db/database.py

@@ -1,12 +1,13 @@
 from google.cloud import firestore
 import logging
+from typing import List, Optional
 
 db = firestore.Client()
 
 logger = logging.getLogger(__name__)
 
 
-def get_user_query_stream(user):
+def get_user_query_stream(user: str) -> List[firestore.DocumentSnapshot]:
 
     users = [i for i in db.collection(u'spotify_users').where(u'username', u'==', user).stream()]
 
@@ -17,7 +18,7 @@ def get_user_query_stream(user):
         return []
 
 
-def get_user_doc_ref(user):
+def get_user_doc_ref(user: str) -> Optional[firestore.DocumentReference]:
 
     users = get_user_query_stream(user)
 
@@ -34,14 +35,14 @@ def get_user_doc_ref(user):
         return None
 
 
-def get_user_playlists_collection(user_id):
+def get_user_playlists_collection(user_id: str) -> firestore.CollectionReference:
 
     playlists = db.document(u'spotify_users/{}'.format(user_id)).collection(u'playlists')
 
     return playlists
 
 
-def get_user_playlist_ref_by_username(user, playlist):
+def get_user_playlist_ref_by_username(user: str, playlist: str) -> Optional[firestore.CollectionReference]:
 
     user_ref = get_user_doc_ref(user)
 
@@ -54,7 +55,8 @@ def get_user_playlist_ref_by_username(user, playlist):
         return None
 
 
-def get_user_playlist_ref_by_user_ref(user_ref, playlist):
+def get_user_playlist_ref_by_user_ref(user_ref: firestore.DocumentReference,
+                                      playlist: str) -> Optional[firestore.CollectionReference]:
 
     playlist_collection = get_user_playlists_collection(user_ref.id)