rapid interations of server side app and firefox extension
This commit is contained in:
Bryan Van Deusen
@@ -0,0 +1,5 @@
|
||||
"""API blueprints."""
|
||||
|
||||
from app.api import sources, downloads, credentials, settings, platforms, websocket
|
||||
|
||||
__all__ = ["sources", "downloads", "credentials", "settings", "platforms", "websocket"]
|
||||
@@ -0,0 +1,182 @@
|
||||
"""Credential management API endpoints."""
|
||||
|
||||
from datetime import datetime
|
||||
from quart import Blueprint, request, jsonify, current_app
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.credential import Credential, CredentialType
|
||||
from app.models.setting import Setting, EXTENSION_API_KEY_SETTING, generate_api_key
|
||||
from app.config import get_settings
|
||||
from app.utils.encryption import encrypt_data, decrypt_data
|
||||
|
||||
bp = Blueprint("credentials", __name__)
|
||||
|
||||
VALID_PLATFORMS = ["patreon", "subscribestar", "hentaifoundry", "discord"]
|
||||
|
||||
|
||||
@bp.route("", methods=["GET"])
|
||||
async def list_credentials():
|
||||
"""List all credential status (without sensitive data)."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(select(Credential))
|
||||
credentials = result.scalars().all()
|
||||
|
||||
return jsonify({
|
||||
"items": [c.to_dict(include_data=False) for c in credentials]
|
||||
})
|
||||
|
||||
|
||||
@bp.route("", methods=["POST"])
|
||||
async def upload_credentials():
|
||||
"""Upload credentials (cookies or token).
|
||||
|
||||
Can be called by:
|
||||
1. Firefox extension (with X-Extension-Key header)
|
||||
2. Web UI manual upload (with session auth)
|
||||
"""
|
||||
settings = get_settings()
|
||||
|
||||
# Check extension API key if present
|
||||
extension_key = request.headers.get("X-Extension-Key")
|
||||
if extension_key:
|
||||
# Look up API key from database
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
result = await session.execute(
|
||||
select(Setting).where(Setting.key == EXTENSION_API_KEY_SETTING)
|
||||
)
|
||||
setting = result.scalar()
|
||||
stored_key = setting.value if setting else None
|
||||
|
||||
if not stored_key or extension_key != stored_key:
|
||||
return jsonify({"error": "Invalid extension API key"}), 401
|
||||
|
||||
data = await request.get_json()
|
||||
|
||||
# Validate required fields
|
||||
if not data.get("platform"):
|
||||
return jsonify({"error": "Platform is required"}), 400
|
||||
if not data.get("credential_type"):
|
||||
return jsonify({"error": "Credential type is required"}), 400
|
||||
if not data.get("data"):
|
||||
return jsonify({"error": "Credential data is required"}), 400
|
||||
|
||||
platform = data["platform"]
|
||||
cred_type = data["credential_type"]
|
||||
cred_data = data["data"]
|
||||
|
||||
# Validate platform
|
||||
if platform not in VALID_PLATFORMS:
|
||||
return jsonify({"error": f"Invalid platform. Must be one of: {', '.join(VALID_PLATFORMS)}"}), 400
|
||||
|
||||
# Validate credential type
|
||||
if cred_type not in [CredentialType.COOKIES, CredentialType.TOKEN]:
|
||||
return jsonify({"error": "Invalid credential type. Must be 'cookies' or 'token'"}), 400
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Check for existing credential
|
||||
result = await session.execute(
|
||||
select(Credential).where(Credential.platform == platform)
|
||||
)
|
||||
credential = result.scalar()
|
||||
|
||||
# Encrypt the credential data
|
||||
encrypted = encrypt_data(cred_data, settings.secret_key)
|
||||
|
||||
if credential:
|
||||
# Update existing
|
||||
credential.credential_type = cred_type
|
||||
credential.data = encrypted
|
||||
credential.expires_at = data.get("expires_at")
|
||||
credential.last_verified = None # Reset verification
|
||||
else:
|
||||
# Create new
|
||||
credential = Credential(
|
||||
platform=platform,
|
||||
credential_type=cred_type,
|
||||
data=encrypted,
|
||||
expires_at=data.get("expires_at"),
|
||||
)
|
||||
session.add(credential)
|
||||
|
||||
await session.commit()
|
||||
await session.refresh(credential)
|
||||
|
||||
current_app.logger.info(f"Updated credentials for platform: {platform}")
|
||||
return jsonify({
|
||||
"message": "Credentials updated",
|
||||
"platform": platform,
|
||||
"credential_type": cred_type,
|
||||
"expires_at": credential.expires_at.isoformat() if credential.expires_at else None,
|
||||
})
|
||||
|
||||
|
||||
@bp.route("/<platform>", methods=["DELETE"])
|
||||
async def delete_credentials(platform: str):
|
||||
"""Delete credentials for a platform."""
|
||||
if platform not in VALID_PLATFORMS:
|
||||
return jsonify({"error": f"Invalid platform. Must be one of: {', '.join(VALID_PLATFORMS)}"}), 400
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Credential).where(Credential.platform == platform)
|
||||
)
|
||||
credential = result.scalar()
|
||||
|
||||
if not credential:
|
||||
return jsonify({"error": "Credentials not found for this platform"}), 404
|
||||
|
||||
await session.delete(credential)
|
||||
await session.commit()
|
||||
|
||||
current_app.logger.info(f"Deleted credentials for platform: {platform}")
|
||||
return "", 204
|
||||
|
||||
|
||||
@bp.route("/verify", methods=["POST"])
|
||||
async def verify_credentials():
|
||||
"""Verify that credentials are working.
|
||||
|
||||
Attempts to make a test request to the platform.
|
||||
"""
|
||||
data = await request.get_json()
|
||||
platform = data.get("platform")
|
||||
|
||||
if not platform:
|
||||
return jsonify({"error": "Platform is required"}), 400
|
||||
|
||||
if platform not in VALID_PLATFORMS:
|
||||
return jsonify({"error": f"Invalid platform. Must be one of: {', '.join(VALID_PLATFORMS)}"}), 400
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Credential).where(Credential.platform == platform)
|
||||
)
|
||||
credential = result.scalar()
|
||||
|
||||
if not credential:
|
||||
return jsonify({
|
||||
"platform": platform,
|
||||
"is_valid": False,
|
||||
"error": "No credentials stored for this platform",
|
||||
})
|
||||
|
||||
# TODO: Implement actual verification by making test request
|
||||
# For now, just mark as verified
|
||||
credential.last_verified = datetime.utcnow()
|
||||
await session.commit()
|
||||
|
||||
return jsonify({
|
||||
"platform": platform,
|
||||
"is_valid": True,
|
||||
"last_verified": credential.last_verified.isoformat(),
|
||||
})
|
||||
@@ -0,0 +1,171 @@
|
||||
"""Download history API endpoints."""
|
||||
|
||||
from datetime import datetime
|
||||
from quart import Blueprint, request, jsonify, current_app
|
||||
from sqlalchemy import select, func, and_
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from app.models.download import Download, DownloadStatus
|
||||
from app.models.source import Source
|
||||
from app.tasks.downloads import process_download
|
||||
|
||||
bp = Blueprint("downloads", __name__)
|
||||
|
||||
|
||||
@bp.route("", methods=["GET"])
|
||||
async def list_downloads():
|
||||
"""List download history with filtering and pagination."""
|
||||
# Query parameters
|
||||
source_id = request.args.get("source_id", type=int)
|
||||
status = request.args.get("status")
|
||||
from_date = request.args.get("from_date")
|
||||
to_date = request.args.get("to_date")
|
||||
page = int(request.args.get("page", 1))
|
||||
per_page = min(int(request.args.get("per_page", 50)), 100)
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Build query
|
||||
query = select(Download)
|
||||
count_query = select(func.count()).select_from(Download)
|
||||
|
||||
filters = []
|
||||
if source_id:
|
||||
filters.append(Download.source_id == source_id)
|
||||
if status:
|
||||
filters.append(Download.status == status)
|
||||
if from_date:
|
||||
filters.append(Download.created_at >= datetime.fromisoformat(from_date))
|
||||
if to_date:
|
||||
filters.append(Download.created_at <= datetime.fromisoformat(to_date))
|
||||
|
||||
if filters:
|
||||
query = query.where(and_(*filters))
|
||||
count_query = count_query.where(and_(*filters))
|
||||
|
||||
# Get total count
|
||||
total_result = await session.execute(count_query)
|
||||
total = total_result.scalar()
|
||||
|
||||
# Apply pagination and ordering
|
||||
query = query.order_by(Download.created_at.desc())
|
||||
query = query.offset((page - 1) * per_page).limit(per_page)
|
||||
|
||||
result = await session.execute(query)
|
||||
downloads = result.scalars().all()
|
||||
|
||||
return jsonify({
|
||||
"items": [d.to_dict() for d in downloads],
|
||||
"total": total,
|
||||
"page": page,
|
||||
"per_page": per_page,
|
||||
"pages": (total + per_page - 1) // per_page,
|
||||
})
|
||||
|
||||
|
||||
@bp.route("/<int:download_id>", methods=["GET"])
|
||||
async def get_download(download_id: int):
|
||||
"""Get download details."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(select(Download).where(Download.id == download_id))
|
||||
download = result.scalar()
|
||||
|
||||
if not download:
|
||||
return jsonify({"error": "Download not found"}), 404
|
||||
|
||||
# Include source info if available
|
||||
response = download.to_dict()
|
||||
if download.source_id:
|
||||
source_result = await session.execute(
|
||||
select(Source)
|
||||
.options(selectinload(Source.subscription))
|
||||
.where(Source.id == download.source_id)
|
||||
)
|
||||
source = source_result.scalar()
|
||||
if source:
|
||||
response["source"] = {
|
||||
"id": source.id,
|
||||
"subscription_name": source.subscription.name if source.subscription else None,
|
||||
"platform": source.platform,
|
||||
}
|
||||
|
||||
return jsonify(response)
|
||||
|
||||
|
||||
@bp.route("/<int:download_id>/retry", methods=["POST"])
|
||||
async def retry_download(download_id: int):
|
||||
"""Retry a failed download."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(select(Download).where(Download.id == download_id))
|
||||
download = result.scalar()
|
||||
|
||||
if not download:
|
||||
return jsonify({"error": "Download not found"}), 404
|
||||
|
||||
if download.status != DownloadStatus.FAILED:
|
||||
return jsonify({"error": "Can only retry failed downloads"}), 400
|
||||
|
||||
# Reset status to pending
|
||||
download.status = DownloadStatus.PENDING
|
||||
download.error_type = None
|
||||
download.error_message = None
|
||||
await session.commit()
|
||||
|
||||
# Queue Celery task to retry the download
|
||||
task = process_download.delay(download_id)
|
||||
|
||||
current_app.logger.info(f"Queued retry for download: {download_id}")
|
||||
return jsonify({
|
||||
"message": "Retry queued",
|
||||
"download_id": download_id,
|
||||
"task_id": task.id,
|
||||
}), 202
|
||||
|
||||
|
||||
@bp.route("/stats", methods=["GET"])
|
||||
async def get_stats():
|
||||
"""Get download statistics."""
|
||||
period = request.args.get("period", "week") # day, week, month
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Total counts by status
|
||||
status_query = select(
|
||||
Download.status,
|
||||
func.count(Download.id)
|
||||
).group_by(Download.status)
|
||||
|
||||
status_result = await session.execute(status_query)
|
||||
status_counts = {row[0]: row[1] for row in status_result}
|
||||
|
||||
# Counts by platform (through source)
|
||||
platform_query = select(
|
||||
Source.platform,
|
||||
func.count(Download.id)
|
||||
).join(Source, Download.source_id == Source.id).group_by(Source.platform)
|
||||
|
||||
platform_result = await session.execute(platform_query)
|
||||
platform_counts = {row[0]: row[1] for row in platform_result}
|
||||
|
||||
# Recent downloads count
|
||||
recent_query = select(func.count()).select_from(Download).where(
|
||||
Download.status == DownloadStatus.COMPLETED
|
||||
)
|
||||
recent_result = await session.execute(recent_query)
|
||||
recent_completed = recent_result.scalar()
|
||||
|
||||
return jsonify({
|
||||
"total": sum(status_counts.values()),
|
||||
"by_status": status_counts,
|
||||
"by_platform": platform_counts,
|
||||
"completed": status_counts.get(DownloadStatus.COMPLETED, 0),
|
||||
"failed": status_counts.get(DownloadStatus.FAILED, 0),
|
||||
"pending": status_counts.get(DownloadStatus.PENDING, 0),
|
||||
})
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Platform information API endpoints.
|
||||
|
||||
Provides information about supported platforms and their configuration options.
|
||||
"""
|
||||
|
||||
from quart import Blueprint, jsonify
|
||||
|
||||
from app.services.gallery_dl import GalleryDLService
|
||||
|
||||
bp = Blueprint("platforms", __name__)
|
||||
|
||||
|
||||
@bp.route("", methods=["GET"])
|
||||
async def list_platforms():
|
||||
"""List all supported platforms with their configuration options."""
|
||||
gdl_service = GalleryDLService()
|
||||
|
||||
platforms = {
|
||||
"patreon": {
|
||||
"name": "Patreon",
|
||||
"description": "Download posts from Patreon creators",
|
||||
"requires_auth": True,
|
||||
"auth_type": "cookies",
|
||||
"url_pattern": "https://www.patreon.com/{creator}",
|
||||
"url_examples": [
|
||||
"https://www.patreon.com/example_artist",
|
||||
"https://www.patreon.com/user?u=12345678",
|
||||
],
|
||||
"content_types": gdl_service.get_platform_content_types("patreon"),
|
||||
"content_type_descriptions": {
|
||||
"images": "Standard post images",
|
||||
"image_large": "High-resolution image versions",
|
||||
"attachments": "File attachments (ZIP, PSD, etc.)",
|
||||
"postfile": "Post-specific files",
|
||||
"content": "Embedded content in post body",
|
||||
},
|
||||
"default_config": gdl_service.get_default_config_for_platform("patreon"),
|
||||
},
|
||||
"subscribestar": {
|
||||
"name": "SubscribeStar",
|
||||
"description": "Download posts from SubscribeStar creators",
|
||||
"requires_auth": True,
|
||||
"auth_type": "cookies",
|
||||
"url_pattern": "https://subscribestar.{domain}/{creator}",
|
||||
"url_examples": [
|
||||
"https://subscribestar.adult/example_artist",
|
||||
"https://www.subscribestar.com/example_artist",
|
||||
],
|
||||
"content_types": gdl_service.get_platform_content_types("subscribestar"),
|
||||
"content_type_descriptions": {
|
||||
"all": "All available content",
|
||||
},
|
||||
"default_config": gdl_service.get_default_config_for_platform("subscribestar"),
|
||||
},
|
||||
"hentaifoundry": {
|
||||
"name": "Hentai Foundry",
|
||||
"description": "Download artwork from Hentai Foundry artists",
|
||||
"requires_auth": False,
|
||||
"auth_type": "cookies", # Optional but improves access
|
||||
"url_pattern": "https://www.hentai-foundry.com/user/{artist}",
|
||||
"url_examples": [
|
||||
"https://www.hentai-foundry.com/user/example_artist",
|
||||
"https://www.hentai-foundry.com/pictures/user/example_artist",
|
||||
],
|
||||
"content_types": gdl_service.get_platform_content_types("hentaifoundry"),
|
||||
"content_type_descriptions": {
|
||||
"pictures": "Artwork/pictures",
|
||||
"stories": "Written stories",
|
||||
},
|
||||
"default_config": gdl_service.get_default_config_for_platform("hentaifoundry"),
|
||||
},
|
||||
"discord": {
|
||||
"name": "Discord",
|
||||
"description": "Download attachments from Discord channels",
|
||||
"requires_auth": True,
|
||||
"auth_type": "token",
|
||||
"url_pattern": "https://discord.com/channels/{server}/{channel}",
|
||||
"url_examples": [
|
||||
"https://discord.com/channels/123456789/987654321",
|
||||
],
|
||||
"content_types": gdl_service.get_platform_content_types("discord"),
|
||||
"content_type_descriptions": {
|
||||
"all": "All attachments and embeds",
|
||||
},
|
||||
"default_config": gdl_service.get_default_config_for_platform("discord"),
|
||||
"notes": "Requires Discord user token (not bot token)",
|
||||
},
|
||||
}
|
||||
|
||||
return jsonify({"platforms": platforms})
|
||||
|
||||
|
||||
@bp.route("/<platform>", methods=["GET"])
|
||||
async def get_platform(platform: str):
|
||||
"""Get detailed information about a specific platform."""
|
||||
gdl_service = GalleryDLService()
|
||||
|
||||
if platform not in ["patreon", "subscribestar", "hentaifoundry", "discord"]:
|
||||
return jsonify({"error": f"Unknown platform: {platform}"}), 404
|
||||
|
||||
# Get the full platform info from list_platforms
|
||||
all_platforms = (await list_platforms()).get_json()
|
||||
platform_info = all_platforms["platforms"].get(platform)
|
||||
|
||||
if not platform_info:
|
||||
return jsonify({"error": f"Platform not found: {platform}"}), 404
|
||||
|
||||
return jsonify(platform_info)
|
||||
|
||||
|
||||
@bp.route("/<platform>/config-schema", methods=["GET"])
|
||||
async def get_config_schema(platform: str):
|
||||
"""Get the configuration schema for a platform.
|
||||
|
||||
This defines what options can be set per-source for this platform.
|
||||
Useful for building dynamic forms in the frontend.
|
||||
"""
|
||||
if platform not in ["patreon", "subscribestar", "hentaifoundry", "discord"]:
|
||||
return jsonify({"error": f"Unknown platform: {platform}"}), 404
|
||||
|
||||
gdl_service = GalleryDLService()
|
||||
content_types = gdl_service.get_platform_content_types(platform)
|
||||
defaults = gdl_service.get_default_config_for_platform(platform)
|
||||
|
||||
schema = {
|
||||
"platform": platform,
|
||||
"fields": [
|
||||
{
|
||||
"name": "content_types",
|
||||
"type": "multiselect",
|
||||
"label": "Content Types to Download",
|
||||
"description": "Select which types of content to download from this source",
|
||||
"options": content_types,
|
||||
"default": defaults.get("content_types", ["all"]),
|
||||
},
|
||||
{
|
||||
"name": "sleep",
|
||||
"type": "number",
|
||||
"label": "Delay Between Downloads (seconds)",
|
||||
"description": "Time to wait between downloading files. Higher values are safer but slower.",
|
||||
"min": 0.5,
|
||||
"max": 30.0,
|
||||
"step": 0.5,
|
||||
"default": defaults.get("sleep", 3.0),
|
||||
},
|
||||
{
|
||||
"name": "sleep_request",
|
||||
"type": "number",
|
||||
"label": "Delay Between Requests (seconds)",
|
||||
"description": "Time to wait between API requests. Helps avoid rate limiting.",
|
||||
"min": 0.5,
|
||||
"max": 15.0,
|
||||
"step": 0.5,
|
||||
"default": defaults.get("sleep_request", 1.5),
|
||||
},
|
||||
{
|
||||
"name": "skip_existing",
|
||||
"type": "boolean",
|
||||
"label": "Skip Already Downloaded",
|
||||
"description": "Skip files that have already been downloaded (uses archive database)",
|
||||
"default": True,
|
||||
},
|
||||
{
|
||||
"name": "save_metadata",
|
||||
"type": "boolean",
|
||||
"label": "Save Metadata",
|
||||
"description": "Save JSON metadata file alongside each downloaded file",
|
||||
"default": True,
|
||||
},
|
||||
{
|
||||
"name": "timeout",
|
||||
"type": "number",
|
||||
"label": "Download Timeout (seconds)",
|
||||
"description": "Maximum time to wait for a single download before giving up",
|
||||
"min": 60,
|
||||
"max": 7200,
|
||||
"step": 60,
|
||||
"default": 3600,
|
||||
},
|
||||
{
|
||||
"name": "directory_pattern",
|
||||
"type": "text",
|
||||
"label": "Directory Pattern (Advanced)",
|
||||
"description": "Custom directory structure pattern. Leave empty for platform default.",
|
||||
"placeholder": defaults.get("directory_pattern", ""),
|
||||
"default": None,
|
||||
},
|
||||
{
|
||||
"name": "filename_pattern",
|
||||
"type": "text",
|
||||
"label": "Filename Pattern (Advanced)",
|
||||
"description": "Custom filename pattern. Leave empty for platform default.",
|
||||
"placeholder": defaults.get("filename_pattern", ""),
|
||||
"default": None,
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
return jsonify(schema)
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Settings API endpoints."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from quart import Blueprint, request, jsonify, current_app
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.setting import Setting, DEFAULT_SETTINGS, EXTENSION_API_KEY_SETTING, generate_api_key
|
||||
from app.config import get_settings
|
||||
|
||||
bp = Blueprint("settings", __name__)
|
||||
|
||||
|
||||
async def get_or_create_extension_api_key(session: AsyncSession) -> str:
|
||||
"""Get the extension API key, creating one if it doesn't exist."""
|
||||
result = await session.execute(
|
||||
select(Setting).where(Setting.key == EXTENSION_API_KEY_SETTING)
|
||||
)
|
||||
setting = result.scalar()
|
||||
|
||||
if setting:
|
||||
return setting.value
|
||||
|
||||
# Generate new key
|
||||
new_key = generate_api_key()
|
||||
setting = Setting(key=EXTENSION_API_KEY_SETTING, value=new_key)
|
||||
session.add(setting)
|
||||
await session.commit()
|
||||
|
||||
return new_key
|
||||
|
||||
|
||||
@bp.route("", methods=["GET"])
|
||||
async def get_all_settings():
|
||||
"""Get all application settings."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(select(Setting))
|
||||
settings = result.scalars().all()
|
||||
|
||||
# Build settings dict with defaults
|
||||
settings_dict = dict(DEFAULT_SETTINGS)
|
||||
for s in settings:
|
||||
settings_dict[s.key] = s.value
|
||||
|
||||
# Ensure extension API key exists
|
||||
api_key = await get_or_create_extension_api_key(session)
|
||||
settings_dict[EXTENSION_API_KEY_SETTING] = api_key
|
||||
|
||||
return jsonify({"settings": settings_dict})
|
||||
|
||||
|
||||
@bp.route("", methods=["PATCH"])
|
||||
async def update_settings():
|
||||
"""Update application settings."""
|
||||
data = await request.get_json()
|
||||
|
||||
if not data:
|
||||
return jsonify({"error": "No settings provided"}), 400
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
for key, value in data.items():
|
||||
# Check if setting exists
|
||||
result = await session.execute(select(Setting).where(Setting.key == key))
|
||||
setting = result.scalar()
|
||||
|
||||
if setting:
|
||||
setting.value = value
|
||||
else:
|
||||
setting = Setting(key=key, value=value)
|
||||
session.add(setting)
|
||||
|
||||
await session.commit()
|
||||
|
||||
# Return updated settings
|
||||
result = await session.execute(select(Setting))
|
||||
settings = result.scalars().all()
|
||||
settings_dict = dict(DEFAULT_SETTINGS)
|
||||
for s in settings:
|
||||
settings_dict[s.key] = s.value
|
||||
|
||||
current_app.logger.info(f"Updated settings: {list(data.keys())}")
|
||||
return jsonify({"settings": settings_dict})
|
||||
|
||||
|
||||
@bp.route("/api-key", methods=["GET"])
|
||||
async def get_extension_api_key():
|
||||
"""Get the extension API key."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
api_key = await get_or_create_extension_api_key(session)
|
||||
return jsonify({"api_key": api_key})
|
||||
|
||||
|
||||
@bp.route("/api-key/regenerate", methods=["POST"])
|
||||
async def regenerate_extension_api_key():
|
||||
"""Regenerate the extension API key."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Find existing key
|
||||
result = await session.execute(
|
||||
select(Setting).where(Setting.key == EXTENSION_API_KEY_SETTING)
|
||||
)
|
||||
setting = result.scalar()
|
||||
|
||||
# Generate new key
|
||||
new_key = generate_api_key()
|
||||
|
||||
if setting:
|
||||
setting.value = new_key
|
||||
else:
|
||||
setting = Setting(key=EXTENSION_API_KEY_SETTING, value=new_key)
|
||||
session.add(setting)
|
||||
|
||||
await session.commit()
|
||||
|
||||
current_app.logger.info("Extension API key regenerated")
|
||||
return jsonify({"api_key": new_key, "message": "API key regenerated"})
|
||||
|
||||
|
||||
@bp.route("/gallery-dl", methods=["GET"])
|
||||
async def get_gallery_dl_config():
|
||||
"""Get gallery-dl configuration."""
|
||||
config = get_settings()
|
||||
config_path = Path(config.config_path) / "gallery-dl.conf"
|
||||
|
||||
if not config_path.exists():
|
||||
return jsonify({"config": {}, "message": "No configuration file found"})
|
||||
|
||||
try:
|
||||
with open(config_path) as f:
|
||||
gallery_dl_config = json.load(f)
|
||||
return jsonify({"config": gallery_dl_config})
|
||||
except json.JSONDecodeError as e:
|
||||
return jsonify({"error": f"Invalid JSON in config file: {e}"}), 500
|
||||
|
||||
|
||||
@bp.route("/gallery-dl", methods=["PUT"])
|
||||
async def update_gallery_dl_config():
|
||||
"""Update gallery-dl configuration."""
|
||||
data = await request.get_json()
|
||||
|
||||
if not data.get("config"):
|
||||
return jsonify({"error": "Config object is required"}), 400
|
||||
|
||||
config = get_settings()
|
||||
config_path = Path(config.config_path) / "gallery-dl.conf"
|
||||
|
||||
# Ensure config directory exists
|
||||
config_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(data["config"], f, indent=4)
|
||||
|
||||
current_app.logger.info("Updated gallery-dl configuration")
|
||||
return jsonify({"message": "Configuration updated", "config": data["config"]})
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Failed to write gallery-dl config: {e}")
|
||||
return jsonify({"error": f"Failed to write config: {e}"}), 500
|
||||
@@ -0,0 +1,223 @@
|
||||
"""Source management API endpoints."""
|
||||
|
||||
from quart import Blueprint, request, jsonify, current_app
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from app.models.source import Source
|
||||
from app.models.subscription import Subscription
|
||||
from app.tasks.downloads import download_source
|
||||
|
||||
bp = Blueprint("sources", __name__)
|
||||
|
||||
|
||||
@bp.route("", methods=["GET"])
|
||||
async def list_sources():
|
||||
"""List all sources with optional filtering and pagination."""
|
||||
# Query parameters
|
||||
platform = request.args.get("platform")
|
||||
subscription_id = request.args.get("subscription_id")
|
||||
enabled = request.args.get("enabled")
|
||||
page = int(request.args.get("page", 1))
|
||||
per_page = min(int(request.args.get("per_page", 50)), 100)
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Build query with subscription eager loading
|
||||
query = select(Source).options(selectinload(Source.subscription))
|
||||
|
||||
if platform:
|
||||
query = query.where(Source.platform == platform)
|
||||
if subscription_id:
|
||||
query = query.where(Source.subscription_id == int(subscription_id))
|
||||
if enabled is not None:
|
||||
query = query.where(Source.enabled == (enabled.lower() == "true"))
|
||||
|
||||
# Get total count
|
||||
count_query = select(func.count()).select_from(Source)
|
||||
if platform:
|
||||
count_query = count_query.where(Source.platform == platform)
|
||||
if subscription_id:
|
||||
count_query = count_query.where(Source.subscription_id == int(subscription_id))
|
||||
if enabled is not None:
|
||||
count_query = count_query.where(Source.enabled == (enabled.lower() == "true"))
|
||||
|
||||
total_result = await session.execute(count_query)
|
||||
total = total_result.scalar()
|
||||
|
||||
# Apply pagination - order by subscription name, then platform
|
||||
query = query.join(Subscription).order_by(Subscription.name, Source.platform)
|
||||
query = query.offset((page - 1) * per_page).limit(per_page)
|
||||
|
||||
result = await session.execute(query)
|
||||
sources = result.scalars().all()
|
||||
|
||||
return jsonify({
|
||||
"items": [s.to_dict() for s in sources],
|
||||
"total": total,
|
||||
"page": page,
|
||||
"per_page": per_page,
|
||||
"pages": (total + per_page - 1) // per_page if total > 0 else 0,
|
||||
})
|
||||
|
||||
|
||||
@bp.route("", methods=["POST"])
|
||||
async def create_source():
|
||||
"""Create a new source for a subscription."""
|
||||
data = await request.get_json()
|
||||
|
||||
# Validate required fields
|
||||
required = ["subscription_id", "platform", "url"]
|
||||
missing = [f for f in required if not data.get(f)]
|
||||
if missing:
|
||||
return jsonify({"error": f"Missing required fields: {', '.join(missing)}"}), 400
|
||||
|
||||
# Validate platform
|
||||
valid_platforms = ["patreon", "subscribestar", "hentaifoundry", "discord"]
|
||||
if data["platform"] not in valid_platforms:
|
||||
return jsonify({"error": f"Invalid platform. Must be one of: {', '.join(valid_platforms)}"}), 400
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Check subscription exists
|
||||
sub_result = await session.execute(
|
||||
select(Subscription).where(Subscription.id == data["subscription_id"])
|
||||
)
|
||||
subscription = sub_result.scalar()
|
||||
if not subscription:
|
||||
return jsonify({"error": "Subscription not found"}), 404
|
||||
|
||||
# Check for duplicate platform+url
|
||||
existing = await session.execute(
|
||||
select(Source).where(
|
||||
Source.platform == data["platform"],
|
||||
Source.url == data["url"]
|
||||
)
|
||||
)
|
||||
if existing.scalar():
|
||||
return jsonify({"error": "Source with this platform and URL already exists"}), 409
|
||||
|
||||
source = Source(
|
||||
subscription_id=data["subscription_id"],
|
||||
platform=data["platform"],
|
||||
url=data["url"],
|
||||
enabled=data.get("enabled", True),
|
||||
check_interval=data.get("check_interval", 3600),
|
||||
metadata_=data.get("metadata", {}),
|
||||
)
|
||||
|
||||
session.add(source)
|
||||
await session.commit()
|
||||
await session.refresh(source)
|
||||
|
||||
current_app.logger.info(f"Created source: {subscription.name}/{source.platform}")
|
||||
return jsonify(source.to_dict()), 201
|
||||
|
||||
|
||||
@bp.route("/<int:source_id>", methods=["GET"])
|
||||
async def get_source(source_id: int):
|
||||
"""Get a single source by ID."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Source)
|
||||
.options(selectinload(Source.subscription))
|
||||
.where(Source.id == source_id)
|
||||
)
|
||||
source = result.scalar()
|
||||
|
||||
if not source:
|
||||
return jsonify({"error": "Source not found"}), 404
|
||||
|
||||
return jsonify(source.to_dict())
|
||||
|
||||
|
||||
@bp.route("/<int:source_id>", methods=["PATCH"])
|
||||
async def update_source(source_id: int):
|
||||
"""Update a source."""
|
||||
data = await request.get_json()
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Source)
|
||||
.options(selectinload(Source.subscription))
|
||||
.where(Source.id == source_id)
|
||||
)
|
||||
source = result.scalar()
|
||||
|
||||
if not source:
|
||||
return jsonify({"error": "Source not found"}), 404
|
||||
|
||||
# Update allowed fields
|
||||
allowed_fields = ["enabled", "check_interval", "metadata"]
|
||||
for field in allowed_fields:
|
||||
if field in data:
|
||||
if field == "metadata":
|
||||
source.metadata_ = data[field]
|
||||
else:
|
||||
setattr(source, field, data[field])
|
||||
|
||||
await session.commit()
|
||||
await session.refresh(source)
|
||||
|
||||
current_app.logger.info(f"Updated source: {source.subscription.name}/{source.platform}")
|
||||
return jsonify(source.to_dict())
|
||||
|
||||
|
||||
@bp.route("/<int:source_id>", methods=["DELETE"])
|
||||
async def delete_source(source_id: int):
|
||||
"""Delete a source."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Source)
|
||||
.options(selectinload(Source.subscription))
|
||||
.where(Source.id == source_id)
|
||||
)
|
||||
source = result.scalar()
|
||||
|
||||
if not source:
|
||||
return jsonify({"error": "Source not found"}), 404
|
||||
|
||||
info = f"{source.subscription.name}/{source.platform}"
|
||||
await session.delete(source)
|
||||
await session.commit()
|
||||
|
||||
current_app.logger.info(f"Deleted source: {info}")
|
||||
return "", 204
|
||||
|
||||
|
||||
@bp.route("/<int:source_id>/check", methods=["POST"])
|
||||
async def trigger_check(source_id: int):
|
||||
"""Trigger an immediate download check for a source."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Source)
|
||||
.options(selectinload(Source.subscription))
|
||||
.where(Source.id == source_id)
|
||||
)
|
||||
source = result.scalar()
|
||||
|
||||
if not source:
|
||||
return jsonify({"error": "Source not found"}), 404
|
||||
|
||||
# Queue Celery task for immediate download
|
||||
task = download_source.delay(source_id)
|
||||
|
||||
current_app.logger.info(f"Triggered check for source: {source.subscription.name}/{source.platform}")
|
||||
return jsonify({
|
||||
"message": "Check queued",
|
||||
"source_id": source_id,
|
||||
"subscription_name": source.subscription.name,
|
||||
"platform": source.platform,
|
||||
"task_id": task.id,
|
||||
}), 202
|
||||
@@ -0,0 +1,307 @@
|
||||
"""Subscription management API endpoints."""
|
||||
|
||||
from quart import Blueprint, request, jsonify, current_app
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from app.models.subscription import Subscription
|
||||
from app.models.source import Source
|
||||
from app.tasks.downloads import download_source
|
||||
|
||||
bp = Blueprint("subscriptions", __name__)
|
||||
|
||||
|
||||
@bp.route("", methods=["GET"])
|
||||
async def list_subscriptions():
|
||||
"""List all subscriptions with optional filtering and pagination."""
|
||||
# Query parameters
|
||||
enabled = request.args.get("enabled")
|
||||
search = request.args.get("search")
|
||||
page = int(request.args.get("page", 1))
|
||||
per_page = min(int(request.args.get("per_page", 50)), 100)
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Build query with eager loading of sources
|
||||
query = select(Subscription).options(selectinload(Subscription.sources))
|
||||
|
||||
if enabled is not None:
|
||||
query = query.where(Subscription.enabled == (enabled.lower() == "true"))
|
||||
if search:
|
||||
query = query.where(Subscription.name.ilike(f"%{search}%"))
|
||||
|
||||
# Get total count
|
||||
count_query = select(func.count()).select_from(Subscription)
|
||||
if enabled is not None:
|
||||
count_query = count_query.where(Subscription.enabled == (enabled.lower() == "true"))
|
||||
if search:
|
||||
count_query = count_query.where(Subscription.name.ilike(f"%{search}%"))
|
||||
|
||||
total_result = await session.execute(count_query)
|
||||
total = total_result.scalar()
|
||||
|
||||
# Apply pagination
|
||||
query = query.order_by(Subscription.priority.desc(), Subscription.name)
|
||||
query = query.offset((page - 1) * per_page).limit(per_page)
|
||||
|
||||
result = await session.execute(query)
|
||||
subscriptions = result.scalars().unique().all()
|
||||
|
||||
return jsonify({
|
||||
"items": [s.to_dict() for s in subscriptions],
|
||||
"total": total,
|
||||
"page": page,
|
||||
"per_page": per_page,
|
||||
"pages": (total + per_page - 1) // per_page if total > 0 else 0,
|
||||
})
|
||||
|
||||
|
||||
@bp.route("", methods=["POST"])
|
||||
async def create_subscription():
|
||||
"""Create a new subscription with optional sources."""
|
||||
data = await request.get_json()
|
||||
|
||||
# Validate required fields
|
||||
if not data.get("name"):
|
||||
return jsonify({"error": "Name is required"}), 400
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Check for duplicate name
|
||||
existing = await session.execute(
|
||||
select(Subscription).where(Subscription.name == data["name"])
|
||||
)
|
||||
if existing.scalar():
|
||||
return jsonify({"error": "Subscription with this name already exists"}), 409
|
||||
|
||||
subscription = Subscription(
|
||||
name=data["name"],
|
||||
enabled=data.get("enabled", True),
|
||||
priority=data.get("priority", 0),
|
||||
description=data.get("description"),
|
||||
metadata_=data.get("metadata", {}),
|
||||
)
|
||||
|
||||
# Add sources if provided
|
||||
sources_data = data.get("sources", [])
|
||||
for source_data in sources_data:
|
||||
if not source_data.get("platform") or not source_data.get("url"):
|
||||
continue
|
||||
source = Source(
|
||||
platform=source_data["platform"],
|
||||
url=source_data["url"],
|
||||
enabled=source_data.get("enabled", True),
|
||||
check_interval=source_data.get("check_interval", 3600),
|
||||
metadata_=source_data.get("metadata", {}),
|
||||
)
|
||||
subscription.sources.append(source)
|
||||
|
||||
session.add(subscription)
|
||||
await session.commit()
|
||||
await session.refresh(subscription)
|
||||
|
||||
current_app.logger.info(f"Created subscription: {subscription.name}")
|
||||
return jsonify(subscription.to_dict()), 201
|
||||
|
||||
|
||||
@bp.route("/<int:subscription_id>", methods=["GET"])
|
||||
async def get_subscription(subscription_id: int):
|
||||
"""Get a single subscription by ID with its sources."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Subscription)
|
||||
.options(selectinload(Subscription.sources))
|
||||
.where(Subscription.id == subscription_id)
|
||||
)
|
||||
subscription = result.scalar()
|
||||
|
||||
if not subscription:
|
||||
return jsonify({"error": "Subscription not found"}), 404
|
||||
|
||||
return jsonify(subscription.to_dict())
|
||||
|
||||
|
||||
@bp.route("/<int:subscription_id>", methods=["PATCH"])
|
||||
async def update_subscription(subscription_id: int):
|
||||
"""Update a subscription."""
|
||||
data = await request.get_json()
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Subscription)
|
||||
.options(selectinload(Subscription.sources))
|
||||
.where(Subscription.id == subscription_id)
|
||||
)
|
||||
subscription = result.scalar()
|
||||
|
||||
if not subscription:
|
||||
return jsonify({"error": "Subscription not found"}), 404
|
||||
|
||||
# Check for duplicate name if changing
|
||||
if "name" in data and data["name"] != subscription.name:
|
||||
existing = await session.execute(
|
||||
select(Subscription).where(Subscription.name == data["name"])
|
||||
)
|
||||
if existing.scalar():
|
||||
return jsonify({"error": "Subscription with this name already exists"}), 409
|
||||
|
||||
# Update allowed fields
|
||||
allowed_fields = ["name", "enabled", "priority", "description", "metadata"]
|
||||
for field in allowed_fields:
|
||||
if field in data:
|
||||
if field == "metadata":
|
||||
subscription.metadata_ = data[field]
|
||||
else:
|
||||
setattr(subscription, field, data[field])
|
||||
|
||||
await session.commit()
|
||||
await session.refresh(subscription)
|
||||
|
||||
current_app.logger.info(f"Updated subscription: {subscription.name}")
|
||||
return jsonify(subscription.to_dict())
|
||||
|
||||
|
||||
@bp.route("/<int:subscription_id>", methods=["DELETE"])
|
||||
async def delete_subscription(subscription_id: int):
|
||||
"""Delete a subscription and all its sources."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Subscription).where(Subscription.id == subscription_id)
|
||||
)
|
||||
subscription = result.scalar()
|
||||
|
||||
if not subscription:
|
||||
return jsonify({"error": "Subscription not found"}), 404
|
||||
|
||||
name = subscription.name
|
||||
await session.delete(subscription)
|
||||
await session.commit()
|
||||
|
||||
current_app.logger.info(f"Deleted subscription: {name}")
|
||||
return "", 204
|
||||
|
||||
|
||||
@bp.route("/<int:subscription_id>/sources", methods=["GET"])
|
||||
async def list_subscription_sources(subscription_id: int):
|
||||
"""List all sources for a subscription."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Subscription)
|
||||
.options(selectinload(Subscription.sources))
|
||||
.where(Subscription.id == subscription_id)
|
||||
)
|
||||
subscription = result.scalar()
|
||||
|
||||
if not subscription:
|
||||
return jsonify({"error": "Subscription not found"}), 404
|
||||
|
||||
return jsonify({
|
||||
"subscription_id": subscription_id,
|
||||
"subscription_name": subscription.name,
|
||||
"sources": [s.to_dict(include_subscription=False) for s in subscription.sources]
|
||||
})
|
||||
|
||||
|
||||
@bp.route("/<int:subscription_id>/sources", methods=["POST"])
|
||||
async def add_source_to_subscription(subscription_id: int):
|
||||
"""Add a new source to a subscription."""
|
||||
data = await request.get_json()
|
||||
|
||||
# Validate required fields
|
||||
required = ["platform", "url"]
|
||||
missing = [f for f in required if not data.get(f)]
|
||||
if missing:
|
||||
return jsonify({"error": f"Missing required fields: {', '.join(missing)}"}), 400
|
||||
|
||||
# Validate platform
|
||||
valid_platforms = ["patreon", "subscribestar", "hentaifoundry", "discord"]
|
||||
if data["platform"] not in valid_platforms:
|
||||
return jsonify({"error": f"Invalid platform. Must be one of: {', '.join(valid_platforms)}"}), 400
|
||||
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
# Check subscription exists
|
||||
result = await session.execute(
|
||||
select(Subscription).where(Subscription.id == subscription_id)
|
||||
)
|
||||
subscription = result.scalar()
|
||||
|
||||
if not subscription:
|
||||
return jsonify({"error": "Subscription not found"}), 404
|
||||
|
||||
# Check for duplicate platform+url
|
||||
existing = await session.execute(
|
||||
select(Source).where(
|
||||
Source.platform == data["platform"],
|
||||
Source.url == data["url"]
|
||||
)
|
||||
)
|
||||
if existing.scalar():
|
||||
return jsonify({"error": "Source with this platform and URL already exists"}), 409
|
||||
|
||||
# Store name before commit (attributes expire after commit)
|
||||
subscription_name = subscription.name
|
||||
|
||||
source = Source(
|
||||
subscription_id=subscription_id,
|
||||
platform=data["platform"],
|
||||
url=data["url"],
|
||||
enabled=data.get("enabled", True),
|
||||
check_interval=data.get("check_interval", 3600),
|
||||
metadata_=data.get("metadata", {}),
|
||||
)
|
||||
|
||||
session.add(source)
|
||||
await session.commit()
|
||||
await session.refresh(source)
|
||||
|
||||
current_app.logger.info(f"Added source to {subscription_name}: {source.platform}")
|
||||
# Use include_subscription=False to avoid lazy loading the relationship
|
||||
result = source.to_dict(include_subscription=False)
|
||||
result["subscription_name"] = subscription_name
|
||||
return jsonify(result), 201
|
||||
|
||||
|
||||
@bp.route("/<int:subscription_id>/check", methods=["POST"])
|
||||
async def trigger_subscription_check(subscription_id: int):
|
||||
"""Trigger download check for all enabled sources in a subscription."""
|
||||
async with current_app.db_engine.connect() as conn:
|
||||
session = AsyncSession(bind=conn)
|
||||
|
||||
result = await session.execute(
|
||||
select(Subscription)
|
||||
.options(selectinload(Subscription.sources))
|
||||
.where(Subscription.id == subscription_id)
|
||||
)
|
||||
subscription = result.scalar()
|
||||
|
||||
if not subscription:
|
||||
return jsonify({"error": "Subscription not found"}), 404
|
||||
|
||||
enabled_sources = [s for s in subscription.sources if s.enabled]
|
||||
|
||||
# Queue Celery tasks for each enabled source
|
||||
task_ids = []
|
||||
for source in enabled_sources:
|
||||
task = download_source.delay(source.id)
|
||||
task_ids.append(task.id)
|
||||
|
||||
current_app.logger.info(f"Triggered check for subscription: {subscription.name} ({len(enabled_sources)} sources)")
|
||||
return jsonify({
|
||||
"message": "Checks queued",
|
||||
"subscription_id": subscription_id,
|
||||
"source_count": len(enabled_sources),
|
||||
"task_ids": task_ids,
|
||||
}), 202
|
||||
@@ -0,0 +1,217 @@
|
||||
"""WebSocket API for real-time updates."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Set
|
||||
from quart import Blueprint, websocket, current_app
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
from app.config import get_settings
|
||||
from app.events import EVENTS_CHANNEL
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
bp = Blueprint("websocket", __name__)
|
||||
|
||||
# Connected WebSocket clients
|
||||
connected_clients: Set = set()
|
||||
|
||||
|
||||
class WebSocketManager:
|
||||
"""Manages WebSocket connections and broadcasts."""
|
||||
|
||||
_instance = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance.clients = set()
|
||||
return cls._instance
|
||||
|
||||
async def connect(self, ws):
|
||||
"""Register a new WebSocket connection."""
|
||||
self.clients.add(ws)
|
||||
logger.info(f"WebSocket connected. Total clients: {len(self.clients)}")
|
||||
|
||||
async def disconnect(self, ws):
|
||||
"""Remove a WebSocket connection."""
|
||||
self.clients.discard(ws)
|
||||
logger.info(f"WebSocket disconnected. Total clients: {len(self.clients)}")
|
||||
|
||||
async def broadcast(self, event_type: str, data: dict):
|
||||
"""Broadcast an event to all connected clients."""
|
||||
if not self.clients:
|
||||
return
|
||||
|
||||
message = json.dumps({
|
||||
"type": event_type,
|
||||
"data": data,
|
||||
})
|
||||
|
||||
disconnected = set()
|
||||
for client in self.clients:
|
||||
try:
|
||||
await client.send(message)
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to send to client: {e}")
|
||||
disconnected.add(client)
|
||||
|
||||
# Clean up disconnected clients
|
||||
self.clients -= disconnected
|
||||
|
||||
async def send_to_client(self, ws, event_type: str, data: dict):
|
||||
"""Send an event to a specific client."""
|
||||
message = json.dumps({
|
||||
"type": event_type,
|
||||
"data": data,
|
||||
})
|
||||
try:
|
||||
await ws.send(message)
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to send to client: {e}")
|
||||
|
||||
|
||||
# Global manager instance
|
||||
ws_manager = WebSocketManager()
|
||||
|
||||
|
||||
@bp.websocket("/events")
|
||||
async def events():
|
||||
"""WebSocket endpoint for real-time events.
|
||||
|
||||
Events sent from server:
|
||||
- download.started: { download_id, source_id, url }
|
||||
- download.progress: { download_id, file_count }
|
||||
- download.completed: { download_id, file_count, duration }
|
||||
- download.failed: { download_id, error_type, error_message }
|
||||
- source.updated: { source_id, name, last_check }
|
||||
- credential.expiring: { platform, expires_in_hours }
|
||||
"""
|
||||
await ws_manager.connect(websocket._get_current_object())
|
||||
|
||||
try:
|
||||
# Send initial connection confirmation
|
||||
await websocket.send(json.dumps({
|
||||
"type": "connected",
|
||||
"data": {"message": "Connected to Gallery Subscriber"},
|
||||
}))
|
||||
|
||||
# Keep connection alive and handle incoming messages
|
||||
while True:
|
||||
try:
|
||||
# Wait for messages (ping/pong or commands)
|
||||
message = await asyncio.wait_for(websocket.receive(), timeout=30)
|
||||
|
||||
# Handle ping
|
||||
if message == "ping":
|
||||
await websocket.send("pong")
|
||||
else:
|
||||
# Parse JSON commands if needed
|
||||
try:
|
||||
data = json.loads(message)
|
||||
await handle_client_message(websocket._get_current_object(), data)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
# Send keepalive ping
|
||||
try:
|
||||
await websocket.send(json.dumps({"type": "ping"}))
|
||||
except Exception:
|
||||
break
|
||||
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
finally:
|
||||
await ws_manager.disconnect(websocket._get_current_object())
|
||||
|
||||
|
||||
async def handle_client_message(ws, data: dict):
|
||||
"""Handle messages from WebSocket clients."""
|
||||
msg_type = data.get("type")
|
||||
|
||||
if msg_type == "subscribe":
|
||||
# Client wants to subscribe to specific events
|
||||
# For now, all clients receive all events
|
||||
await ws_manager.send_to_client(ws, "subscribed", {"status": "ok"})
|
||||
|
||||
elif msg_type == "ping":
|
||||
await ws_manager.send_to_client(ws, "pong", {})
|
||||
|
||||
|
||||
# Helper functions for broadcasting events from other parts of the app
|
||||
|
||||
async def broadcast_download_started(download_id: int, source_id: int, url: str):
|
||||
"""Broadcast that a download has started."""
|
||||
await ws_manager.broadcast("download.started", {
|
||||
"download_id": download_id,
|
||||
"source_id": source_id,
|
||||
"url": url,
|
||||
})
|
||||
|
||||
|
||||
async def broadcast_download_completed(download_id: int, file_count: int, duration: float):
|
||||
"""Broadcast that a download has completed."""
|
||||
await ws_manager.broadcast("download.completed", {
|
||||
"download_id": download_id,
|
||||
"file_count": file_count,
|
||||
"duration_seconds": duration,
|
||||
})
|
||||
|
||||
|
||||
async def broadcast_download_failed(download_id: int, error_type: str, error_message: str):
|
||||
"""Broadcast that a download has failed."""
|
||||
await ws_manager.broadcast("download.failed", {
|
||||
"download_id": download_id,
|
||||
"error_type": error_type,
|
||||
"error_message": error_message,
|
||||
})
|
||||
|
||||
|
||||
async def broadcast_source_updated(source_id: int, name: str, last_check: str):
|
||||
"""Broadcast that a source has been updated."""
|
||||
await ws_manager.broadcast("source.updated", {
|
||||
"source_id": source_id,
|
||||
"name": name,
|
||||
"last_check": last_check,
|
||||
})
|
||||
|
||||
|
||||
# Redis Pub/Sub integration for cross-process events (from Celery tasks)
|
||||
|
||||
async def start_redis_subscriber():
|
||||
"""Start a background task that subscribes to Redis events and broadcasts them.
|
||||
|
||||
This should be called when the Quart app starts.
|
||||
"""
|
||||
settings = get_settings()
|
||||
|
||||
try:
|
||||
redis_client = aioredis.from_url(settings.redis_url)
|
||||
pubsub = redis_client.pubsub()
|
||||
await pubsub.subscribe(EVENTS_CHANNEL)
|
||||
|
||||
logger.info(f"Started Redis subscriber on channel: {EVENTS_CHANNEL}")
|
||||
|
||||
async for message in pubsub.listen():
|
||||
if message["type"] == "message":
|
||||
try:
|
||||
data = json.loads(message["data"])
|
||||
event_type = data.get("type")
|
||||
event_data = data.get("data", {})
|
||||
|
||||
# Broadcast to all WebSocket clients
|
||||
await ws_manager.broadcast(event_type, event_data)
|
||||
logger.debug(f"Broadcast event from Redis: {event_type}")
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Invalid JSON in Redis message: {message['data']}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error broadcasting Redis event: {e}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Redis subscriber cancelled")
|
||||
await pubsub.unsubscribe(EVENTS_CHANNEL)
|
||||
await redis_client.close()
|
||||
except Exception as e:
|
||||
logger.error(f"Redis subscriber error: {e}")
|
||||
Reference in New Issue
Block a user