From 00353837b432f0b59761dd8e9d269bdb95d8018d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= <andras@schmelczer.dev>
Date: Wed, 25 May 2022 19:37:11 +0200
Subject: [PATCH] Improve CLI DX

---
 examples/{main_service.py => main.py}     |  2 +-
 great_ai/src/great_ai/__main__.py         | 33 +++++++++++++++++------
 great_ai/src/great_ai/open_s3/__main__.py |  3 +++
 3 files changed, 29 insertions(+), 9 deletions(-)
 rename examples/{main_service.py => main.py} (77%)

diff --git a/examples/main_service.py b/examples/main.py
similarity index 77%
rename from examples/main_service.py
rename to examples/main.py
index 29049d5..20d5537 100755
--- a/examples/main_service.py
+++ b/examples/main.py
@@ -2,7 +2,7 @@
 
 from great_ai import configure, create_service
 
-configure(development_mode_override=True)
+# configure(development_mode_override=True)
 
 from predict_domain import predict_domain
 
diff --git a/great_ai/src/great_ai/__main__.py b/great_ai/src/great_ai/__main__.py
index fd9eb0e..29a1bdb 100644
--- a/great_ai/src/great_ai/__main__.py
+++ b/great_ai/src/great_ai/__main__.py
@@ -2,28 +2,37 @@
 import re
 from importlib import import_module
 from sys import argv
-
 import uvicorn
 from uvicorn.config import LOGGING_CONFIG
 
-from great_ai.great_ai.context import get_context
+from .great_ai.exceptions import MissingArgumentError
+from .great_ai.context import get_context
 
-if __name__ == "__main__":
+
+
+def main():
     if len(argv) < 2:
-        raise ValueError(f"Provide a filename such as: {argv[0]} my_app.py")
+        raise MissingArgumentError(f"Provide a filename such as: {argv[0]} my_app.py")
 
     module_name = re.sub(r"\.py\b", "", argv[1])
+    
+    base = module_name.split(":")[0]
+
+    import_module(base)
+    get_context().logger.info("Starting server")
+
     if ":" not in module_name:
         module_name += ":app"
-
-    base = module_name.split(":")[0]
-    module = import_module(base)
-    get_context().logger.info("Starting server")
+        get_context().logger.warning(
+            'Service name (name of variable) is assumed to be "app",'
+            + ' such as: `app = create_service(predict_domain)`'
+        )
 
     uvicorn.run(
         module_name,
         host="0.0.0.0",
         port=6060,
+        timeout_keep_alive=600,
         workers=1,
         reload=not get_context().is_production,
         log_config={
@@ -40,3 +49,11 @@ if __name__ == "__main__":
             },
         },
     )
+
+if __name__ == "__main__":
+    try:
+        main()
+    except (MissingArgumentError, ModuleNotFoundError) as e:
+        get_context().logger.error(e)
+    except KeyboardInterrupt:
+        exit()
diff --git a/great_ai/src/great_ai/open_s3/__main__.py b/great_ai/src/great_ai/open_s3/__main__.py
index cfc46ec..901f8c9 100644
--- a/great_ai/src/great_ai/open_s3/__main__.py
+++ b/great_ai/src/great_ai/open_s3/__main__.py
@@ -32,5 +32,8 @@ if __name__ == "__main__":
         if args.delete:
             for f in args.delete:
                 LargeFile(f).delete()
+    except KeyboardInterrupt:
+        logger.warning("Exiting")
+        exit()
     except Exception as e:
         logger.error(e)