clean CSV code

ckanext/tracking/blueprints/csv.py

@@ -17,14 +17,12 @@
 tracking_csv_blueprint = Blueprint('tracking_csv', __name__, url_prefix='/tracking-csv')
 
 
-@tracking_csv_blueprint.route('/most-accessed-dataset-with-token.csv', methods=["GET"])
-def most_accessed_dataset_with_token_csv():
-    """ Get most accessed (using a API token) datasets """
-
+def _csv_response(auth_fn, data_fn, kwargs, filename):
+    """ general CSV response from data_fn(**kwargs) checking access with auth_fn """
     current_user_name = current_user.name if current_user else None
     context = {'user': current_user_name}
-    toolkit.check_access('most_accessed_dataset_with_token_csv', context)
-    data = most_accessed_dataset_with_token_data(limit=10)
+    toolkit.check_access(auth_fn, context)
+    data = data_fn(**kwargs)
     # If no rows, return empty 204 CSV
     if len(data) == 0:
         return Response('', status=204)
@@ -38,58 +36,40 @@ def most_accessed_dataset_with_token_csv():
         writer.writerow(row)
 
     response = Response(buffer.getvalue(), mimetype='text/csv')
-    filename = 'most-accessed-dataset-with-token.csv'
     response.headers.set("Content-Disposition", "attachment", filename=filename)
     return response
 
 
-@tracking_csv_blueprint.route('/most-accessed-token.csv', methods=["GET"])
-def most_accessed_token_csv():
-    """ Get most accessed tokens """
+@tracking_csv_blueprint.route('/most-accessed-dataset-with-token.csv', methods=["GET"])
+def most_accessed_dataset_with_token_csv():
+    """ Get most accessed (using a API token) datasets """
 
-    current_user_name = current_user.name if current_user else None
-    context = {'user': current_user_name}
-    toolkit.check_access('most_accessed_token_csv', context)
-    data = most_accessed_token_data(limit=10)
-    # If no rows, return empty 204 CSV
-    if len(data) == 0:
-        return Response('', status=204)
-    headers = data[0].keys()
-    buffer = StringIO()
-    writer = csv.DictWriter(buffer, fieldnames=headers)
+    return _csv_response(
+        'most_accessed_dataset_with_token_csv',
+        most_accessed_dataset_with_token_data,
+        {'limit': 10},
+        'most-accessed-dataset-with-token.csv',
+    )
 
-    writer.writeheader()
-    for row in data:
-        writer.writerow(row)
 
-    response = Response(buffer.getvalue(), mimetype='text/csv')
-    filename = 'most-accessed-tokens.csv'
-    response.headers.set("Content-Disposition", "attachment", filename=filename)
-    return response
+@tracking_csv_blueprint.route('/most-accessed-token.csv', methods=["GET"])
+def most_accessed_token_csv():
+    """ Get most accessed tokens """
+    return _csv_response(
+        'most_accessed_token_csv',
+        most_accessed_token_data,
+        {'limit': 10},
+        'most-accessed-tokens.csv',
+    )
 
 
 @tracking_csv_blueprint.route('/all-token-usage.csv', methods=["GET"])
 def all_token_usage_csv():
     """ Get all tokens usage """
 
-    current_user_name = current_user.name if current_user else None
-    context = {'user': current_user_name}
-    toolkit.check_access('most_accessed_token_csv', context)
-    data = all_token_usage_data(limit=1000)
-    # If no rows, return empty 204 CSV
-    if len(data) == 0:
-        return Response('', status=204)
-
-    # Create CSV including package details
-    headers = data[0].keys()
-    buffer = StringIO()
-    writer = csv.DictWriter(buffer, fieldnames=headers)
-
-    writer.writeheader()
-    for row in data:
-        writer.writerow(row)
-
-    response = Response(buffer.getvalue(), mimetype='text/csv')
-    filename = 'all-token-usage.csv'
-    response.headers.set("Content-Disposition", "attachment", filename=filename)
-    return response
+    return _csv_response(
+        'all_token_usage_csv',
+        all_token_usage_data,
+        {'limit': 1000},
+        'all-token-usage.csv',
+    )