LB-1532: Add API endpoints for user data export

listenbrainz/webserver/__init__.py

@@ -243,7 +243,7 @@ def create_web_app(debug=None):
     htmx = HTMX(app)
 
     # Static files
-    import listenbrainz.webserver.static_manager
+    import listenbrainz.webserver.static_manager as static_manager
     static_manager.read_manifest()
     app.static_folder = '/static'
 
@@ -457,3 +457,6 @@ def _register_blueprints(app):
 
     from listenbrainz.webserver.views.internet_archive_api import internet_archive_api_bp
     app.register_blueprint(internet_archive_api_bp, url_prefix=API_PREFIX+"/internet_archive")
+
+    from listenbrainz.webserver.views.export_api import export_api_bp
+    app.register_blueprint(export_api_bp, url_prefix=API_PREFIX+'/export')

listenbrainz/webserver/views/export_api.py

@@ -0,0 +1,130 @@
+import json
+import os
+from flask import Blueprint, current_app, jsonify, send_file
+from sqlalchemy import text
+from psycopg2 import DatabaseError
+
+from listenbrainz.webserver import db_conn
+from listenbrainz.webserver.decorators import api_listenstore_needed
+from listenbrainz.webserver.errors import APIBadRequest, APIInternalServerError, APINotFound
+from listenbrainz.webserver.views.api_tools import validate_auth_header
+
+export_api_bp = Blueprint("export_api", __name__)
+
+
+@export_api_bp.route("/", methods=["POST"])
+@api_listenstore_needed
+def create_export_task():
+    """ Add a request to export the user data to an archive in background. """
+    user = validate_auth_header()
+    user_id = user["id"]
+    try:
+        query = """
+            INSERT INTO user_data_export (user_id, type, status, progress)
+                 VALUES (:user_id, :type, 'waiting', :progress)
+            ON CONFLICT (user_id, type)
+                  WHERE status = 'waiting' OR status = 'in_progress'
+             DO NOTHING
+              RETURNING id, type, available_until, created, progress, status, filename
+        """
+        result = db_conn.execute(text(query), {
+            "user_id": user_id,
+            "type": "export_all_user_data",
+            "progress": "Your data export will start soon."
+        })
+        export = result.first()
+
+        if export is not None:
+            query = "INSERT INTO background_tasks (user_id, task, metadata) VALUES (:user_id, :task, :metadata) ON CONFLICT DO NOTHING RETURNING id"
+            result = db_conn.execute(text(query), {
+                "user_id": user_id,
+                "task": "export_all_user_data",
+                "metadata": json.dumps({"export_id": export.id})
+            })
+            task = result.first()
+            if task is not None:
+                db_conn.commit()
+                return jsonify({
+                    "export_id": export.id,
+                    "type": export.type,
+                    "available_until": export.available_until.isoformat() if export.available_until is not None else None,
+                    "created": export.created.isoformat(),
+                    "progress": export.progress,
+                    "status": export.status,
+                    "filename": export.filename,
+                })
+
+        # task already exists in queue, rollback new entry
+        db_conn.rollback()
+        raise APIBadRequest(message="Data export already requested.")
+
+    except DatabaseError:
+        current_app.logger.error('Error while exporting user data: %s', user["musicbrainz_id"], exc_info=True)
+        raise APIInternalServerError(f'Error while exporting user data {user["musicbrainz_id"]}, please try again later.')
+
+
+@export_api_bp.route("/<int:export_id>", methods=["GET"])
+@api_listenstore_needed
+def get_export_task(export_id):
+    """ Retrieve the requested export's data if it belongs to the specified user """
+    user = validate_auth_header()
+    user_id = user["id"]
+
+    result = db_conn.execute(
+        text("SELECT * FROM user_data_export WHERE user_id = :user_id AND id = :export_id"),
+        {"user_id": user_id, "export_id": export_id}
+    )
+    row = result.first()
+    if row is None:
+        raise APINotFound("Export not found")
+    return jsonify({
+        "export_id": row.id,
+        "type": row.type,
+        "available_until": row.available_until.isoformat() if row.available_until is not None else None,
+        "created": row.created.isoformat(),
+        "progress": row.progress,
+        "status": row.status,
+        "filename": row.filename,
+    })
+
+
+@export_api_bp.route("/", methods=["GET"])
+@api_listenstore_needed
+def list_export_tasks():
+    """ Retrieve the all export tasks for the current user """
+    user = validate_auth_header()
+    user_id = user["id"]
+
+    result = db_conn.execute(
+        text("SELECT * FROM user_data_export WHERE user_id = :user_id ORDER BY created DESC"),
+        {"user_id": user_id}
+    )
+    rows = result.mappings().all()
+    return jsonify([{
+        "export_id": row["id"],
+        "type": row["type"],
+        "available_until": row["available_until"].isoformat() if row["available_until"] is not None else None,
+        "created": row["created"].isoformat(),
+        "progress": row["progress"],
+        "status": row["status"],
+        "filename": row["filename"],
+    } for row in rows])
+
+
+@export_api_bp.route("/<int:export_id>/download", methods=["GET"])
+@api_listenstore_needed
+def download_export_archive(export_id):
+    """ Download the requested export if it is complete and belongs to the specified user """
+    user = validate_auth_header()
+    user_id = user["id"]
+
+    result = db_conn.execute(
+        text("SELECT filename FROM user_data_export WHERE user_id = :user_id AND status = 'completed' AND id = :export_id"),
+        {"user_id": user_id, "export_id": export_id}
+    )
+    row = result.first()
+    if row is None:
+        raise APINotFound("Export not found")
+
+    file_path = os.path.join(current_app.config["USER_DATA_EXPORT_BASE_DIR"], str(row["filename"]))
+    return send_file(file_path, mimetype="application/zip", as_attachment=True)

listenbrainz/webserver/views/test/test_export_api.py

@@ -0,0 +1,70 @@
+import unittest
+from flask import json
+from listenbrainz.tests.integration import ListenAPIIntegrationTestCase
+from listenbrainz.db.user import get_or_create
+
+
+class ExportAPIIntegrationTestCase(ListenAPIIntegrationTestCase):
+
+    def setUp(self):
+        super().setUp()
+        # Ensure we have a user with a token
+        self.user = get_or_create(self.db_conn, 1, 'testuser')
+        self.auth_headers = {'Authorization': f'Token {self.user["auth_token"]}'}
+
+    def test_create_get_list_export(self):
+        """ Test the complete flow of creating, checking, and listing exports """
+
+        # 1. Create an export task
+        response = self.client.post('/1/export/', headers=self.auth_headers)
+        self.assertEqual(response.status_code, 200)
+        data = json.loads(response.data)
+
+        self.assertIn('export_id', data)
+        self.assertEqual(data['status'], 'waiting')
+        self.assertEqual(data['type'], 'export_all_user_data')
+
+        export_id = data['export_id']
+
+        # 2. Check the status of the export
+        response = self.client.get(f'/1/export/{export_id}', headers=self.auth_headers)
+        self.assertEqual(response.status_code, 200)
+        data = json.loads(response.data)
+
+        self.assertEqual(data['export_id'], export_id)
+        # In a test environment without a real background worker, status remains 'waiting'
+        self.assertEqual(data['status'], 'waiting')
+
+        # 3. List all exports
+        response = self.client.get('/1/export/', headers=self.auth_headers)
+        self.assertEqual(response.status_code, 200)
+        data = json.loads(response.data)
+
+        self.assertIsInstance(data, list)
+        self.assertGreaterEqual(len(data), 1)
+
+        # Verify our export is in the list
+        found = False
+        for export in data:
+            if export['export_id'] == export_id:
+                found = True
+                break
+        self.assertTrue(found, "Created export not found in list")
+
+    def test_export_unauthorized(self):
+        """ Test that endpoints require authentication """
+        response = self.client.post('/1/export/')
+        self.assertEqual(response.status_code, 401)
+
+        response = self.client.get('/1/export/')
+        self.assertEqual(response.status_code, 401)
+
+    def test_get_nonexistent_export(self):
+        """ Test getting a non-existent export """
+        # Assuming 999999 is an ID that doesn't exist
+        response = self.client.get('/1/export/999999', headers=self.auth_headers)
+        self.assertEqual(response.status_code, 404)
+
+
+if __name__ == '__main__':
+    unittest.main()