From bd590c15c93a775aa464ba8b4057e915a7f22276 Mon Sep 17 00:00:00 2001
From: Avram Lubkin <aviso@rockhopper.net>
Date: Tue, 11 Jul 2023 10:28:51 -0400
Subject: [PATCH] Consolidate direct program exits

---
 comma/cli/__init__.py     | 13 +++++++++----
 comma/database/driver.py  |  6 +++---
 comma/exceptions.py       | 21 +++++++++++++++++++++
 comma/util/spreadsheet.py | 13 +++++++------
 4 files changed, 40 insertions(+), 13 deletions(-)
 create mode 100644 comma/exceptions.py

diff --git a/comma/cli/__init__.py b/comma/cli/__init__.py
index daad20a..8878b4f 100644
--- a/comma/cli/__init__.py
+++ b/comma/cli/__init__.py
@@ -17,6 +17,7 @@ from comma.config import BasicConfig, FullConfig
 from comma.database.driver import DatabaseDriver
 from comma.database.model import Distros, MonitoringSubjects
 from comma.downstream import Downstream
+from comma.exceptions import CommaError
 from comma.upstream import Upstream
 from comma.util.spreadsheet import Spreadsheet
 from comma.util.symbols import Symbols
@@ -172,11 +173,15 @@ def main(args: Optional[Sequence[str]] = None):
     else:
         config: BasicConfig = BasicConfig(**vars(options))
 
-    # Get database object
-    database = DatabaseDriver(dry_run=options.dry_run, echo=options.verbose > 2)
+    try:
+        # Get database object
+        database = DatabaseDriver(dry_run=options.dry_run, echo=options.verbose > 2)
 
-    # Create session object and invoke subcommand
-    Session(config, database)(options)
+        # Create session object and invoke subcommand
+        Session(config, database)(options)
+
+    except CommaError as e:
+        sys.exit(f"ERROR: {e}")
 
 
 if __name__ == "__main__":
diff --git a/comma/database/driver.py b/comma/database/driver.py
index 71c0531..4108b4d 100755
--- a/comma/database/driver.py
+++ b/comma/database/driver.py
@@ -7,13 +7,13 @@ Provide a class for managing database connections and sessions
 
 import logging
 import os
-import sys
 import urllib
 from contextlib import contextmanager
 
 import sqlalchemy
 
 from comma.database.model import Base, Distros, MonitoringSubjects
+from comma.exceptions import CommaDatabaseError, CommaDataError
 
 
 LOGGER = logging.getLogger(__name__)
@@ -49,7 +49,7 @@ class DatabaseDriver:
         # Verify credentials are available
         for envvar in ("COMMA_DB_URL", "COMMA_DB_NAME", "COMMA_DB_USERNAME", "COMMA_DB_PW"):
             if not os.environ.get(envvar):
-                sys.exit(f"{envvar} is not defined in the current environment")
+                raise CommaDatabaseError(f"{envvar} is not defined in the current environment")
 
         params = urllib.parse.quote_plus(
             ";".join(
@@ -137,7 +137,7 @@ class DatabaseDriver:
 
             # If URL wasn't given, make sure repo is in database
             elif (name,) not in session.query(Distros.distroID).all():
-                sys.exit(f"Repository '{name}' given without URL not found in database")
+                raise CommaDataError(f"Repository '{name}' given without URL not found in database")
 
             # Add target
             session.add(MonitoringSubjects(distroID=name, revision=revision))
diff --git a/comma/exceptions.py b/comma/exceptions.py
new file mode 100644
index 0000000..ddb20bc
--- /dev/null
+++ b/comma/exceptions.py
@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Custom exception types
+"""
+
+
+class CommaError(Exception):
+    """Comma parent exception class"""
+
+
+class CommaDatabaseError(CommaError):
+    """Errors establishing database connections"""
+
+
+class CommaDataError(CommaError):
+    """Errors with input or stored data"""
+
+
+class CommaSpreadsheetError(CommaError):
+    """Errors with spreadsheet operations"""
diff --git a/comma/util/spreadsheet.py b/comma/util/spreadsheet.py
index a6a7114..fb923fe 100644
--- a/comma/util/spreadsheet.py
+++ b/comma/util/spreadsheet.py
@@ -6,7 +6,6 @@ Functions for exporting data to Excel spreadsheets
 
 import logging
 import re
-import sys
 from datetime import datetime
 from functools import lru_cache
 from pathlib import Path
@@ -21,6 +20,7 @@ from openpyxl.workbook.workbook import Workbook
 from openpyxl.worksheet.worksheet import Worksheet
 
 from comma.database.model import MonitoringSubjects, PatchData
+from comma.exceptions import CommaSpreadsheetError
 
 
 LOGGER = logging.getLogger(__name__)
@@ -113,8 +113,8 @@ def get_workbook(in_file: str) -> Tuple[Workbook, WorksheetWrapper]:
 
     """
     if not Path(in_file).exists():
-        LOGGER.error("The file %s does not exist", in_file)
-        sys.exit(1)
+        raise CommaSpreadsheetError(f"Input spreadsheet '{in_file}' does not exist")
+
     workbook = openpyxl.load_workbook(filename=in_file)
 
     # Force refresh of pivot table in “Pivot” worksheet.
@@ -236,9 +236,10 @@ class Spreadsheet:
                 # Make sure there is a column in the spreadsheet
                 try:
                     worksheet.get_column(repo)
-                except StopIteration:
-                    LOGGER.error("No column with distro '%s', please fix spreadsheet!", repo)
-                    sys.exit(1)
+                except StopIteration as e:
+                    raise CommaSpreadsheetError(
+                        f"No column with distro '{repo}', please fix spreadsheet!"
+                    ) from e
 
                 # Get the latest monitoring subject for the remote
                 targets[repo] = (