From 570f0e3701b67f328a918c636bbcc6d8f5c1951a Mon Sep 17 00:00:00 2001
From: kelle <kellecruz@gmail.com>
Date: Tue, 23 Jul 2024 16:49:43 -0400
Subject: [PATCH] error handling

---
 simple/schema.py               |  3 +--
 simple/utils/spectral_types.py | 49 +++++++++++++++++++---------------
 tests/test_utils.py            |  4 +--
 3 files changed, 29 insertions(+), 27 deletions(-)

diff --git a/simple/schema.py b/simple/schema.py
index 634d73ad3..6461f0d08 100644
--- a/simple/schema.py
+++ b/simple/schema.py
@@ -306,7 +306,6 @@ class SpectralTypes(Base):
     regime = Column(
         String(30),
         ForeignKey("Regimes.regime", ondelete="cascade", onupdate="cascade"),
-        nullable=True,
         primary_key=True,
     )
     adopted = Column(Boolean)
@@ -318,7 +317,7 @@ class SpectralTypes(Base):
         primary_key=True,
     )
 
-    @validates("source", "spectral_type_code", "reference")
+    @validates("source", "spectral_type_code", "regime", "reference")
     def validate_required(self, key, value):
         if value is None:
             raise ValueError(f"Value required for {key}")
diff --git a/simple/utils/spectral_types.py b/simple/utils/spectral_types.py
index 5ba443c18..4dc69e2a3 100644
--- a/simple/utils/spectral_types.py
+++ b/simple/utils/spectral_types.py
@@ -17,6 +17,7 @@
 ]
 
 logger = logging.getLogger("SIMPLE")
+logger.setLevel(logging.INFO)  # set default log level to INFO
 
 
 def ingest_spectral_type(
@@ -72,6 +73,26 @@ def ingest_spectral_type(
         db.query(db.SpectralTypes).filter(db.SpectralTypes.c.source == db_name).table()
     )
 
+    # Check for duplicates
+    duplicate_check = (
+        db.query(db.SpectralTypes.c.source)
+        .filter(
+            and_(
+                db.SpectralTypes.c.source == db_name,
+                db.SpectralTypes.c.regime == regime,
+                db.SpectralTypes.c.reference == reference,
+            )
+        )
+        .count()
+    )
+    if duplicate_check > 0:
+        msg = f"Spectral type already in the database: {db_name}, {regime}, {reference}"
+        if raise_error:
+            logger.error(msg)
+            raise AstroDBError(msg)
+        else:
+            logger.warning(msg)
+
     # set adopted flag
     if source_spt_data is None or len(source_spt_data) == 0:
         adopted = True
@@ -99,9 +120,12 @@ def ingest_spectral_type(
                     adopted = True
                 logger.debug(f"The new spectral type's adopted flag is:, {adopted}")
     else:
-        msg = "Unexpected state"
-        logger.error(msg)
-        raise RuntimeError
+        msg = f"Unexpected state while ingesting {source}"
+        if raise_error:
+            logger.error(msg)
+            raise RuntimeError
+        else:
+            logger.warning(msg)
 
     if spectral_type_code is None:
         spectral_type_code = convert_spt_string_to_code(spectral_type_string)
@@ -121,25 +145,6 @@ def ingest_spectral_type(
         "reference": reference,
     }
 
-    check = (
-        db.query(db.SpectralTypes.c.source)
-        .filter(
-            and_(
-                db.SpectralTypes.c.source == db_name,
-                db.SpectralTypes.c.regime == regime,
-                db.SpectralTypes.c.reference == reference,
-            )
-        )
-        .count()
-    )
-    if check == 1:
-        msg = f"Spectral type for {db_name} already in the database"
-        if raise_error:
-            logger.error(msg)
-            raise AstroDBError(msg)
-        else:
-            logger.warning(msg)
-
     logger.debug(f"Trying to insert {spt_data} into Spectral Types table ")
 
     try:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 4bf6a8da4..bb57a6dba 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -146,9 +146,7 @@ def test_ingest_spectral_type_errors(temp_db):
             reference=spt_data4["reference"],
             regime=spt_data4["regime"],
         )
-    assert "Spectral type for Fake 1 already in the database" in str(
-        error_message.value
-    )
+    assert "Spectral type already in the database" in str(error_message.value)
     # assert "The publication does not exist in the database" in str(error_message.value)