81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
"""Application management CLI: ``python -m photo_pipeline {serve,migrate,worker,import-legacy-scores}``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from typing import Sequence
|
|
|
|
from photo_pipeline.config import Config
|
|
from photo_pipeline.db import run_migrations
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(prog="photo_pipeline")
|
|
commands = parser.add_subparsers(dest="command", required=True)
|
|
commands.add_parser("serve", help="Run the API server")
|
|
commands.add_parser("migrate", help="Upgrade the database to the latest revision")
|
|
worker_cmd = commands.add_parser("worker", help="Run a durable-job worker")
|
|
worker_cmd.add_argument("--id", default="worker-1", help="Worker id (lease owner)")
|
|
import_cmd = commands.add_parser(
|
|
"import-legacy-scores",
|
|
help="Import the archived CLI's nsfw_scores.csv into the database (US07-01)",
|
|
)
|
|
import_cmd.add_argument("csv", help="Path to nsfw_scores.csv")
|
|
import_cmd.add_argument(
|
|
"--overwrite", action="store_true", help="Replace differing imported scores"
|
|
)
|
|
import_cmd.add_argument(
|
|
"--dry-run", action="store_true", help="Report what would happen and change nothing"
|
|
)
|
|
args = parser.parse_args(argv)
|
|
|
|
config = Config.from_env()
|
|
config.database_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
if args.command == "migrate":
|
|
run_migrations(config.database_url)
|
|
return 0
|
|
|
|
if args.command == "import-legacy-scores":
|
|
import json
|
|
|
|
from photo_pipeline.db import create_db_engine, create_session_factory
|
|
from photo_pipeline.services.legacy_import import LegacyImportService, write_report
|
|
|
|
run_migrations(config.database_url)
|
|
engine = create_db_engine(config.database_url)
|
|
service = LegacyImportService(create_session_factory(engine))
|
|
report = service.import_nsfw_scores(
|
|
args.csv, overwrite=args.overwrite, dry_run=args.dry_run
|
|
)
|
|
# The report is the point: an import nobody can audit is not a migration.
|
|
if not args.dry_run:
|
|
write_report(report, config.data_dir)
|
|
print(json.dumps(report.counts, indent=2))
|
|
return 0
|
|
|
|
if args.command == "worker":
|
|
from photo_pipeline.db import create_db_engine, create_session_factory
|
|
|
|
# Importing this registers the safety/analysis job handlers into the shared
|
|
# REGISTRY the Worker defaults to; without it a standalone worker process
|
|
# claims nothing because it knows no job types.
|
|
import photo_pipeline.jobs.domain_handlers # noqa: F401
|
|
from photo_pipeline.jobs.worker import Worker
|
|
|
|
run_migrations(config.database_url)
|
|
engine = create_db_engine(config.database_url)
|
|
Worker(create_session_factory(engine), worker_id=args.id, config=config).run_forever()
|
|
return 0
|
|
|
|
import uvicorn
|
|
|
|
from photo_pipeline.api.app import create_app
|
|
|
|
uvicorn.run(create_app(config), host=config.host, port=config.port)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|