add imdb full import mode

unwind/__main__.py

@@ -1,15 +1,18 @@
+import argparse
 import asyncio
 import logging
+from pathlib import Path
 
 from . import config
 from .db import close_connection_pool, open_connection_pool
 from .imdb import load_imdb
+from .imdb_import import import_from_file
 from .request import session
 
 log = logging.getLogger(__name__)
 
 
-async def run_import():
+async def run_load_user_ratings_from_imdb():
     await open_connection_pool()
 
     with session() as s:
@@ -22,6 +25,60 @@ async def run_import():
     await close_connection_pool()
 
 
+async def run_import_imdb_dataset(basics_path: Path, ratings_path: Path):
+    await open_connection_pool()
+
+    await import_from_file(basics_path=basics_path, ratings_path=ratings_path)
+
+    await close_connection_pool()
+
+
+def getargs():
+    parser = argparse.ArgumentParser()
+    commands = parser.add_subparsers(required=True)
+
+    parser_import_imdb_dataset = commands.add_parser(
+        "import-imdb-dataset",
+        help="Import IMDb datasets.",
+        description="""
+            Import IMDb datasets.
+            New datasets available from https://www.imdb.com/interfaces/.
+        """,
+    )
+    parser_import_imdb_dataset.add_argument(
+        dest="mode",
+        action="store_const",
+        const="import-imdb-dataset",
+    )
+    parser_import_imdb_dataset.add_argument(
+        "--basics", metavar="basics_file.tsv.gz", type=Path, required=True
+    )
+    parser_import_imdb_dataset.add_argument(
+        "--ratings", metavar="ratings_file.tsv.gz", type=Path, required=True
+    )
+
+    parser_load_user_ratings_from_imdb = commands.add_parser(
+        "load-user-ratings-from-imdb",
+        help="Load user ratings from imdb.com.",
+        description="""
+            Refresh user ratings for all registered users live from IMDb's website.
+        """,
+    )
+    parser_load_user_ratings_from_imdb.add_argument(
+        dest="mode",
+        action="store_const",
+        const="load-user-ratings-from-imdb",
+    )
+
+    try:
+        args = parser.parse_args()
+    except TypeError:
+        parser.print_usage()
+        raise
+
+    return args
+
+
 def main():
     logging.basicConfig(
         format="%(asctime)s.%(msecs)03d [%(name)s:%(process)d] %(levelname)s: %(message)s",
@@ -30,7 +87,15 @@ def main():
     )
     log.debug(f"Log level: {config.loglevel}")
 
-    asyncio.run(run_import())
+    try:
+        args = getargs()
+    except:
+        return
+
+    if args.mode == "load-user-ratings-from-imdb":
+        asyncio.run(run_load_user_ratings_from_imdb())
+    elif args.mode == "import-imdb-dataset":
+        asyncio.run(run_import_imdb_dataset(args.basics, args.ratings))
 
 
 main()