diff --git a/run.py b/run.py index 4bfdfa0..58b4c60 100644 --- a/run.py +++ b/run.py @@ -2,20 +2,88 @@ from werkzeug.middleware.proxy_fix import ProxyFix from mirrormanager2.app import create_app from flask_session import Session from cachelib.file import FileSystemCache +from flask import Flask import os +import sys +import subprocess +from threading import Thread -password = os.environ.get('MM2_DATABASE_PASSWORD') -if password: - user = os.environ.get('DB_USER') - host = os.environ.get('DB_HOST') - port = os.environ.get('DB_PORT') - name = os.environ.get('DB_NAME') - os.environ["MM2_SQLALCHEMY_DATABASE_URI"] = f"postgresql://{user}:{password}@{host}:{port}/{name}" -application = create_app() -application.debug = os.environ.get("MM2_DEBUG", False) -application.config['SESSION_TYPE'] = "cachelib" -application.config['SESSION_CACHELIB'] = FileSystemCache(cache_dir='/tmp/sessions', threshold=500) -Session(application) -application.wsgi_app = ProxyFix(application.wsgi_app, x_proto=1, x_host=1) +def setup_env(): + password = os.environ.get('MM2_DATABASE_PASSWORD') + if password: + user = os.environ.get('DB_USER') + host = os.environ.get('DB_HOST') + port = os.environ.get('DB_PORT') + name = os.environ.get('DB_NAME') + os.environ["MM2_SQLALCHEMY_DATABASE_URI"] = f"postgresql://{user}:{password}@{host}:{port}/{name}" + +def mirrormanager_wsgi() -> Flask: + application = create_app() + application.debug = os.environ.get("MM2_DEBUG", False) + application.config['SESSION_TYPE'] = "cachelib" + application.config['SESSION_CACHELIB'] = FileSystemCache(cache_dir='/mnt/efs/fs/0/tmp/sessions', threshold=500) + Session(application) + application.wsgi_app = ProxyFix(application.wsgi_app, x_proto=1, x_host=1) + + +def execute(func: str) -> int: + def stream_output(pipe, output_stream): + while True: + data = pipe.read(1024) # Read in chunks of 1024 bytes + if not data: + break + output_stream.write(data) + output_stream.flush() + pipe.close() + + argv = func.split() + cmd = argv.pop(0) + args = argv if argv else [] + + print(f'running: {cmd} with args: {args}') + + process = subprocess.Popen([cmd] + args, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + bufsize=0) # Unbuffered binary mode + + stdout_thread = Thread(target=stream_output, args=(process.stdout, sys.stdout.buffer)) + stderr_thread = Thread(target=stream_output, args=(process.stderr, sys.stderr.buffer)) + + stdout_thread.start() + stderr_thread.start() + + process.wait() + + stdout_thread.join() + stderr_thread.join() + + return process.returncode + + +if __name__ == "__main__": + setup_env() + # map cmd name to script/function + cmd_map = { + "check-propagation": "/usr/bin/echo undefined", + "crawl-mirrors": "/usr/bin/echo undefined", + "generate-worldmaps": "/opt/app-root/bin/mm2_generate-worldmap", + "scan-primary-mirror": "/usr/local/bin/scan-primary-mirror", + "update-clouds": "/usr/bin/echo undefined", + "update-geoip": "/usr/bin/echo undefined", + "generate-mirrorlist-cache": "/usr/local/bin/generate-mirrorlist-cache", + "update-netblocks": "/opt/app-root/bin/mm2_get-netblocks", + "update-EC2-netblocks": "/opt/app-root/bin/update-EC2-netblocks", + "mirrorlist-server": "/usr/local/bin/mirrorlist-server", + "mirrormanager-wsgi": mirrormanager_wsgi + } + + arg = sys.argv[1] if len(sys.argv) > 1 else "mirrormanager-wsgi" + if arg and (cmd := cmd_map.get(arg, None)): + if not isinstance(cmd, str): + application = cmd() + exit(0) + + sys.exit(execute(cmd))