diff --git a/simple/utils/companions.py b/simple/utils/companions.py index a8ec76c0c..9f93312ba 100644 --- a/simple/utils/companions.py +++ b/simple/utils/companions.py @@ -1,7 +1,7 @@ import logging import sqlalchemy.exc from astrodb_utils import ( - AstroDBError, + AstroDBError, find_source_in_db ) @@ -90,6 +90,14 @@ def ingest_companion_relationships( msg = f"{source}: Source cannot be the same as companion name" logger.error(msg) raise AstroDBError(msg) + + source_name = find_source_in_db(db, source) + if len(source_name) != 1: + msg = f"{source}: No source or multiple sources found: {source_name}" + logger.error(msg) + raise AstroDBError(msg) + else: + source_name = source_name[0] if projected_separation_arcsec is not None and projected_separation_arcsec < 0: msg = f"Projected separation: {projected_separation_arcsec}, cannot be negative" @@ -118,7 +126,7 @@ def ingest_companion_relationships( conn.execute( db.CompanionRelationships.insert().values( { - "source": source, + "source": source_name, "companion_name": companion_name, "projected_separation_arcsec": projected_separation_arcsec, "projected_separation_error": projected_separation_error, @@ -133,7 +141,7 @@ def ingest_companion_relationships( logger.info( "ComapnionRelationship added: ", [ - source, + source_name, companion_name, relationship, projected_separation_arcsec, diff --git a/tests/test_data.py b/tests/test_data.py index f3deffbd0..56610e067 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -627,3 +627,15 @@ def test_radial_velocities(db): .astropy() ) assert len(t) == 89, f"found {len(t)} radial velociies with no uncertainty" + +def test_companion_relations(db): + t = db.query(db.CompanionRelationships).astropy() + assert len(t) == 102, f"found {len(t)} companion relationships" + + ref = "Roth24" + t = ( + db.query(db.CompanionRelationships) + .filter(db.CompanionRelationships.c.reference == ref) + .astropy() + ) + assert len(t) == 89, f"found {len(t)} companion relationships with {ref} reference" \ No newline at end of file