refactor: changed scribe commands + inserting some ZTF exclusive fiel…

…ds in SQL

prv_candidates_step/prv_candidates_step/core/strategy/ztf.py

@@ -23,8 +23,8 @@ def prv_forced_photometry_mapper() -> dict:
 
     mapping.update(
         {
-            "e_ra": Mapper(lambda: 0),
-            "e_dec": Mapper(lambda: 0),
+            "e_ra": Mapper(lambda x: 0),
+            "e_dec": Mapper(lambda x: 0),
             "isdiffpos": Mapper(
                 lambda x: 1 if x >= 0 else -1, origin="forcediffimflux"
             ),
@@ -152,7 +152,7 @@ def extract_detections_and_non_detections(alert: dict) -> dict:
         )
         candidate.update(
             {
-                "pid": fp["pid"],
+                "candid": f'{candidate["oid"]}-{candidate["pid"]}',
                 "aid": aid,
                 "has_stamp": False,  # ?
                 "forced": True,

scribe/mongo_scribe/mongo/command/commands.py

@@ -85,6 +85,8 @@ def _check_inputs(self, collection, data, criteria):
         super()._check_inputs(collection, data, criteria)
         if not criteria:
             raise UpdateWithNoCriteriaException()
+        if collection == "forced_photometry":
+            self.criteria = { "_id": self.criteria["_id"] }
 
     def get_operations(self) -> list:
         op = "$setOnInsert" if self.options.set_on_insert else "$set"

scribe/mongo_scribe/sql/command/commands.py

@@ -122,7 +122,7 @@ class InsertForcedPhotometryCommand(Command):
     type = ValidCommands.insert_forced_photo
 
     def _format_data(self, data: Dict):
-        exclude = ["aid", "sid", "tid", "pid", "e_dec", "e_ra", "stellar", "extra_fields"]
+        exclude = ["aid", "sid", "tid", "e_dec", "e_ra", "stellar", "extra_fields"]
         fid_map = {"g": 1, "r": 2, "i": 3}
 
         extra_fields = data["extra_fields"]
@@ -133,11 +133,11 @@ def _format_data(self, data: Dict):
         new_data = {k: v for k, v in new_data.items() if k not in exclude}
         new_data["fid"] = fid_map[new_data["fid"]]
 
-        return {**new_data, **extra_fields, "candid": self.criteria["candid"]}
+        return {**new_data, **extra_fields}
 
     @staticmethod
     def db_operation(session: Session, data: List):
-        unique = {(el["candid"], el["oid"]): el for el in data}
+        unique = {(el["pid"], el["oid"]): el for el in data}
         unique = list(unique.values())
         return session.connection().execute(
             insert(ForcedPhotometry).values(unique).on_conflict_do_nothing()

scribe/tests/integration/test_step_sql.py

@@ -543,11 +543,12 @@ def test_forced_photometry_insertion(self):
                 "collection": "forced_photometry",
                 "type": "update",
                 "criteria": {
-                    "_id": "candid",
-                    "candid": "12345678-1",
+                    "_id": "candid-pid",
+                    "candid": "candid-pid",
                 },
                 "data": {
                     "mag": 10.0,
+                    "pid": 423432,
                     "oid": "ZTF04ululeea",
                     "corrected": False,
                     "dubious": False,

sorting_hat_step/sorting_hat_step/utils/database.py

@@ -75,13 +75,29 @@ def update_query(db: MongoConnection, records: List[dict]):
 
 def insert_empty_objects_to_sql(db: PsqlConnection, records: List[Dict]):
     # insert into db values = records on conflict do nothing
-    oids = [r["oid"] for r in records if r["sid"].lower() == "ztf"]
-    oids = set(oids)
+    def format_extra_fields(record):
+        extra_fields = record["extra_fields"]
+        return {
+            "ndethist": extra_fields["ndethist"],
+            "ncovhist": extra_fields["ncovhist"],
+            "mjdstarthist": extra_fields["jdstarthist"] - 2400000.5,
+            "mjdendhist": extra_fields["jdendhist"] - 2400000.5,
+        }
+
+    oids = {
+        r["oid"]: format_extra_fields(r) for r in records if r["sid"].lower() == "ztf"
+    }
     with db.session() as session:
-        to_insert = [{"oid": oid} for oid in oids]
+        to_insert = [{"oid": oid, **extra_fields} for oid, extra_fields in oids.items()]
         statement = insert(Object).values(to_insert)
         statement = statement.on_conflict_do_update(
-            "object_pkey", set_=dict(oid=statement.excluded.oid)
+            "object_pkey",
+            set_=dict(
+                ndethist=statement.excluded.ndethist,
+                ncovhist=statement.excluded.ncovhist,
+                mjdstarthist=statement.excluded.mjdstarthist,
+                mjdendhist=statement.excluded.mjdendhist,
+            ),
         )
         session.execute(statement)
         session.commit()

sorting_hat_step/sorting_hat_step/utils/wizard.py

@@ -242,7 +242,7 @@ def insert_empty_objects(mongodb: MongoConnection, alerts: pd.DataFrame, psql=No
     :param db: Connection to the database.
     :alerts: Dataframe with alerts.
     """
-    objects = alerts[["oid", "aid", "sid"]]
+    objects = alerts[["oid", "aid", "sid", "extra_fields"]]
     objects = objects.rename(columns={"aid": "_id"})
     mongo_objects = objects.groupby("_id").oid.apply(list).reset_index()
     logger.debug(

sorting_hat_step/tests/unittest/data/batch.py

@@ -72,6 +72,10 @@ def _generate_ztf_batch(n: int, nearest: int = 0) -> List[dict]:
             "cutoutDifference": {"stampData": b"difference"},
             "candidate": {
                 "jd": random.randrange(2458000, 2459000),
+                "ndethist": random.randint(5,15),
+                "ncovhist": random.randint(5,15),
+                "jdstarthist": random.randint(2458000,2459000),
+                "jdendhist": random.randint(2458000,2459000),
                 "ra": random.uniform(0, 360),
                 "dec": random.uniform(-90, 90),
                 "magpsf": random.uniform(15, 20),