From 6e7b98da845cd7f313835aa0874fa8cfde4f6046 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= <clement.doumouro@gmail.com>
Date: Mon, 10 Jul 2023 11:52:55 +0200
Subject: [PATCH] chore: improve error handling before extension startup

---
 neo4j-app/neo4j_app/core/config.py        |  5 ++-
 neo4j-app/neo4j_app/core/utils/logging.py |  6 ++++
 neo4j-app/neo4j_app/run/run.py            | 43 ++++++++++++++++++-----
 neo4j-app/neo4j_app/tests/run/test_run.py |  2 ++
 4 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/neo4j-app/neo4j_app/core/config.py b/neo4j-app/neo4j_app/core/config.py
index 78264581..726af3fe 100644
--- a/neo4j-app/neo4j_app/core/config.py
+++ b/neo4j-app/neo4j_app/core/config.py
@@ -14,6 +14,7 @@
 
 from neo4j_app.core.elasticsearch import ESClientABC
 from neo4j_app.core.elasticsearch.client import ESClient, OSClient
+from neo4j_app.core.utils.logging import DATE_FMT, STREAM_HANDLER_FMT
 from neo4j_app.core.utils.pydantic import (
     BaseICIJModel,
     IgnoreExtraModel,
@@ -22,8 +23,6 @@
 
 _SYSLOG_MODEL_SPLIT_CHAR = "@"
 _SYSLOG_FMT = f"%(name)s{_SYSLOG_MODEL_SPLIT_CHAR}%(message)s"
-_STREAM_HANDLER_FMT = "[%(levelname)s][%(asctime)s.%(msecs)03d][%(name)s]: %(message)s"
-_DATE_FMT = "%H:%M:%S"
 
 
 def _es_version():
@@ -195,7 +194,7 @@ def setup_loggers(self):
     @functools.cached_property
     def _handlers(self) -> List[logging.Handler]:
         stream_handler = logging.StreamHandler(sys.stderr)
-        stream_handler.setFormatter(logging.Formatter(_STREAM_HANDLER_FMT, _DATE_FMT))
+        stream_handler.setFormatter(logging.Formatter(STREAM_HANDLER_FMT, DATE_FMT))
         handlers = [stream_handler]
         if self.neo4j_app_syslog_facility is not None:
             syslog_handler = SysLogHandler(
diff --git a/neo4j-app/neo4j_app/core/utils/logging.py b/neo4j-app/neo4j_app/core/utils/logging.py
index 1075ef58..97b03b41 100644
--- a/neo4j-app/neo4j_app/core/utils/logging.py
+++ b/neo4j-app/neo4j_app/core/utils/logging.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import logging
 from datetime import datetime
@@ -40,3 +42,7 @@ def log_elapsed_time_cm(
     if "elapsed_time" in output_msg:
         msg_fmt["elapsed_time"] = end
     logger.log(level, output_msg.format(**msg_fmt))
+
+
+STREAM_HANDLER_FMT = "[%(levelname)s][%(asctime)s.%(msecs)03d][%(name)s]: %(message)s"
+DATE_FMT = "%H:%M:%S"
diff --git a/neo4j-app/neo4j_app/run/run.py b/neo4j-app/neo4j_app/run/run.py
index df5540b4..8f828070 100644
--- a/neo4j-app/neo4j_app/run/run.py
+++ b/neo4j-app/neo4j_app/run/run.py
@@ -1,13 +1,17 @@
 # TODO: rename this into run_http ?
 import argparse
+import logging
 import sys
+import traceback
 from pathlib import Path
 from typing import Optional
 
 import uvicorn
 
+import neo4j_app
 from neo4j_app.app.utils import create_app
 from neo4j_app.core.config import AppConfig
+from neo4j_app.core.utils.logging import DATE_FMT, STREAM_HANDLER_FMT
 
 
 def debug_app():
@@ -61,15 +65,38 @@ def get_arg_parser():
     return arg_parser
 
 
-def main():
-    arg_parser = get_arg_parser()
-    args = arg_parser.parse_args()
+def _setup_loggers():
+    loggers = [neo4j_app.__name__, "__main__"]
+    level = logging.INFO
+    stream_handler = logging.StreamHandler(sys.stderr)
+    stream_handler.setFormatter(logging.Formatter(STREAM_HANDLER_FMT, DATE_FMT))
+    for logger in loggers:
+        logger = logging.getLogger(logger)
+        logger.setLevel(level)
+        logger.handlers = []
+        logger.addHandler(stream_handler)
 
-    if hasattr(args, "func"):
-        args.func(args)
-    else:
-        arg_parser.print_help()
-        sys.exit(1)
+
+def main():
+    # Setup loggers temporarily before loggers init using the app configuration
+    _setup_loggers()
+    logger = logging.getLogger(__name__)
+    try:
+        arg_parser = get_arg_parser()
+        args = arg_parser.parse_args()
+
+        if hasattr(args, "func"):
+            args.func(args)
+        else:
+            arg_parser.print_help()
+            sys.exit(1)
+    except KeyboardInterrupt as e:
+        logger.error("Application shutdown...")
+        raise e
+    except Exception as e:  # pylint: disable=broad-except:
+        error_with_trace = "".join(traceback.format_exception(None, e, e.__traceback__))
+        logger.error("Error occurred at application startup:\n%s", error_with_trace)
+        raise e
 
 
 if __name__ == "__main__":
diff --git a/neo4j-app/neo4j_app/tests/run/test_run.py b/neo4j-app/neo4j_app/tests/run/test_run.py
index e50428a8..7b774127 100644
--- a/neo4j-app/neo4j_app/tests/run/test_run.py
+++ b/neo4j-app/neo4j_app/tests/run/test_run.py
@@ -32,4 +32,6 @@ def test_should_read_java_properties(tmpdir: Path):
     # Then
     assert exception.returncode == 1
     assert "Provided config path does not exists" in exception.stderr
+    # Check that the logger was used
+    assert "ERROR" in exception.stderr
     assert str(missing_config_file_path) in exception.stderr