From 344c01f06b8ab2e334782a31cf299cd03406261d Mon Sep 17 00:00:00 2001
From: Kian-Tat Lim <ktl@slac.stanford.edu>
Date: Wed, 1 May 2024 17:44:44 -0700
Subject: [PATCH] Fix update.

---
 python/lsst/consdb/pqclient.py |  8 +++++--
 python/lsst/consdb/pqserver.py | 42 ++++++++++++++++------------------
 2 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/python/lsst/consdb/pqclient.py b/python/lsst/consdb/pqclient.py
index 70688b1a..9ea926ac 100644
--- a/python/lsst/consdb/pqclient.py
+++ b/python/lsst/consdb/pqclient.py
@@ -202,6 +202,7 @@ def insert_flexible_metadata(
         obs_type: str,
         obs_id: int,
         values: dict[str, Any],
+        allow_update: bool = False,
         **kwargs,
     ) -> requests.Response:
         values.update(kwargs)
@@ -214,6 +215,8 @@ def insert_flexible_metadata(
             "obs",
             quote(str(obs_id)),
         )
+        if allow_update:
+            url += "?u=1"
         return self._handle_post(url, data)
 
     def insert(
@@ -227,8 +230,9 @@ def insert(
     ) -> requests.Response:
         values.update(kwargs)
         data = {"table": table, "obs_id": obs_id, "values": values}
-        op = "upsert" if allow_update else "insert"
-        url = urljoin(self.url, op, quote(instrument))
+        url = urljoin(self.url, "insert", quote(instrument), quote(table), "obs", quote(str(obs_id)))
+        if allow_update:
+            url += "?u=1"
         return self._handle_post(url, data)
 
     def query(self, query: str) -> Table:
diff --git a/python/lsst/consdb/pqserver.py b/python/lsst/consdb/pqserver.py
index ab788223..f35ee0ff 100644
--- a/python/lsst/consdb/pqserver.py
+++ b/python/lsst/consdb/pqserver.py
@@ -191,7 +191,7 @@ def get_flexible_metadata(instrument: str, obs_type: str, obs_id: int):
     table = instrument_tables.get_flexible_metadata_table(instrument, obs_type)
     result = dict()
     stmt = sqlalchemy.select(table.c["key", "value"]).where(table.c.obs_id == obs_id)
-    if request.args:
+    if request.args and "k" in request_args:
         cols = request.args["k"]
         stmt = stmt.where(table.c.key in cols)
     logger.debug(str(stmt))
@@ -219,7 +219,16 @@ def insert_flexible_metadata(instrument: str, obs_type: str, obs_id: int):
     elif dtype == "float" and str(float(value)) != value:
         raise BadValueException("float value", value, [])
 
-    stmt = sqlalchemy.insert(table).values(obs_id=obs_id, key=key, value=value)
+    if request.args and request.args.get("u") == "1":
+        stmt = (
+            sqlalchemy.dialects.postgresql.insert(table)
+            .values(obs_id=obs_id, key=key, value=value)
+            .on_conflict_do_update(
+                index_dlements=["obs_id", "key"], set_={"value": value}
+            )
+        )
+    else:
+        stmt = sqlalchemy.insert(table).values(obs_id=obs_id, key=key, value=value)
     logger.debug(str(stmt))
     with engine.connect() as conn:
         _ = conn.execute(stmt)
@@ -228,7 +237,7 @@ def insert_flexible_metadata(instrument: str, obs_type: str, obs_id: int):
 
 
 @app.post("/consdb/insert/<instrument>")
-def insert(instrument: str, upsert: bool = False):
+def insert(instrument: str):
     logger.info(request)
     if instrument not in instrument_tables.schemas:
         raise BadValueException(
@@ -238,27 +247,16 @@ def insert(instrument: str, upsert: bool = False):
     table_name = f"cdb_{instrument}." + info["table"]
     table = instrument_tables.schemas[instrument].tables[table_name]
     valdict = info["values"]
+    valdict["obs_id"] = info["obs_id"]
 
-    stmt = sqlalchemy.insert(table).values(valdict)
-    logger.debug(str(stmt))
-    with engine.connect() as conn:
-        _ = conn.execute(stmt)
-        conn.commit()
-    return ("OK", 200)
-
-
-@app.post("/consdb/upsert/<instrument>")
-def upsert(instrument: str):
-    logger.info(request)
-    if instrument not in instrument_tables.schemas:
-        raise BadValueException(
-            "instrument", instrument, list(instrument_tables.schemas.keys())
+    if request.args and request.args.get("u") == "1":
+        stmt = (
+            sqlalchemy.dialects.postgresql.insert(table)
+            .values(valdict)
+            .on_conflict_do_update(index_elements=["obs_id"], set_=valdict)
         )
-    info = request.json
-    table = instrument + "." + info["table"]
-    valdict = info["values"]
-
-    stmt = sqlalchemy.update(table).values(valdict)
+    else:
+        stmt = sqlalchemy.insert(table).values(valdict)
     logger.debug(str(stmt))
     with engine.connect() as conn:
         _ = conn.execute(stmt)