Downloader/src/main.py

105 lines
3.3 KiB
Python

import asyncio
import os
import signal
import sys
import typing as t
from dataclasses import dataclass
import aiohttp
from flask import Flask, Response, render_template
@dataclass
class Config():
base_url: str
token: t.Optional[str]
user: str
repo: str
task_name: str
page_size: int
artifact_names: t.List[str]
debug: bool
config = {
"base_url": os.environ.get('FORGEJO_BASE_URL'),
"token": os.environ.get('FORGEJO_TOKEN') or None,
"user": os.environ.get('FORGEJO_USER'),
"repo": os.environ.get('FORGEJO_REPO'),
"task_name": os.environ.get('FORGEJO_TASK_NAME'),
"page_size": os.environ.get('FORGEJO_PAGE_SIZE') or 10,
"artifact_names": os.environ.get('FORGEJO_ARTIFACT_NAMES', '').split(','),
"debug": os.environ.get('DEBUG', '') or False,
}
config = Config(**config)
if not config.base_url:
raise ValueError("FORGEJO_BASE_URL is required")
if not config.user:
raise ValueError("FORGEJO_USER is required")
if not config.repo:
raise ValueError("FORGEJO_REPO is required")
if not config.task_name:
raise ValueError("FORGEJO_TASK_NAME is required")
if not config.artifact_names:
raise ValueError("FORGEJO_ARTIFACT_NAMES is required")
app = Flask(__name__, template_folder='templates')
async def fetch_tasks():
url = f'{config.base_url}/api/v1/repos/{config.user}/{config.repo}/actions/tasks'
headers = {
"accept": "application/json",
}
if config.token:
headers["Authorization"] = f"token {config.token}"
params = {
"page_size": config.page_size,
}
async with aiohttp.ClientSession() as session:
async with session.get(url, headers=headers, params=params) as response:
response.raise_for_status()
data = await response.json()
for task in data['workflow_runs']:
if task['name'] == config.task_name:
return task['url']
async def fetch_artifact(artifact_url):
async with aiohttp.ClientSession() as session:
async with session.get(artifact_url) as response:
data = await response.read()
return data
async def handle_artifact(artifact_name: str):
task = await fetch_tasks()
artifact_url = f"{task}/artifacts/{artifact_name}"
artifact_data = await fetch_artifact(artifact_url)
response = Response(artifact_data)
response.headers['Content-Disposition'] = 'attachment; filename="GalacticFactory.zip"'
return response
def create_handle_artifact(artifact_name: str):
async def handle_artifact_async():
return await handle_artifact(artifact_name)
return handle_artifact_async
for artifact_name in config.artifact_names:
route = f'/{artifact_name}/GalacticFactory'
endpoint = f'handle_{artifact_name}'
app.add_url_rule(route, view_func=create_handle_artifact(artifact_name), endpoint=endpoint)
@app.route('/')
def index():
return render_template('index.html', artifact_urls=[f'/{artifact_name}/GalacticFactory' for artifact_name in config.artifact_names])
def handle_sigterm(*args):
print("Received SIGTERM, exiting gracefully...")
sys.exit(0)
signal.signal(signal.SIGTERM, handle_sigterm)
if __name__ == '__main__':
from hypercorn.asyncio import serve
from hypercorn.config import Config as HypercornConfig
h_config = HypercornConfig()
h_config.bind = ["0.0.0.0:80"]
asyncio.run(serve(app=app, config=h_config))